revert: fix(integrations): ensure that GEN_AI_AGENT_NAME is properly set for GEN_AI spans under an invoke_agent span

Author: alexander-alderman-webb

sentry_sdk/integrations/langchain.py

@@ -1,6 +1,4 @@
-import contextvars
 import itertools
-import sys
 import json
 import warnings
 from collections import OrderedDict
@@ -26,10 +24,8 @@
 if TYPE_CHECKING:
     from typing import (
         Any,
-        AsyncIterator,
         Callable,
         Dict,
-        Iterator,
         List,
         Optional,
         Union,
@@ -153,44 +149,6 @@ def _transform_langchain_message_content(content: "Any") -> "Any":
     return content
 
 
-# Contextvar to track agent names in a stack for re-entrant agent support
-_agent_stack: "contextvars.ContextVar[Optional[List[Optional[str]]]]" = (
-    contextvars.ContextVar("langchain_agent_stack", default=None)
-)
-
-
-def _push_agent(agent_name: "Optional[str]") -> None:
-    """Push an agent name onto the stack."""
-    stack = _agent_stack.get()
-    if stack is None:
-        stack = []
-    else:
-        # Copy the list to maintain contextvar isolation across async contexts
-        stack = stack.copy()
-    stack.append(agent_name)
-    _agent_stack.set(stack)
-
-
-def _pop_agent() -> "Optional[str]":
-    """Pop an agent name from the stack and return it."""
-    stack = _agent_stack.get()
-    if stack:
-        # Copy the list to maintain contextvar isolation across async contexts
-        stack = stack.copy()
-        agent_name = stack.pop()
-        _agent_stack.set(stack)
-        return agent_name
-    return None
-
-
-def _get_current_agent() -> "Optional[str]":
-    """Get the current agent name (top of stack) without removing it."""
-    stack = _agent_stack.get()
-    if stack:
-        return stack[-1]
-    return None
-
-
 def _get_system_instructions(messages: "List[List[BaseMessage]]") -> "List[str]":
     system_instructions = []
 
@@ -455,10 +413,6 @@ def on_chat_model_start(
             elif "openai" in ai_type:
                 span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
 
-            agent_name = _get_current_agent()
-            if agent_name:
-                span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
-
             for key, attribute in DATA_FIELDS.items():
                 if key in all_params and all_params[key] is not None:
                     set_data_normalized(span, attribute, all_params[key], unpack=False)
@@ -655,10 +609,6 @@ def on_tool_start(
             if tool_description is not None:
                 span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_description)
 
-            agent_name = _get_current_agent()
-            if agent_name:
-                span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
-
             if should_send_default_pii() and self.include_prompts:
                 set_data_normalized(
                     span,
@@ -985,50 +935,45 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
             name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
             origin=LangchainIntegration.origin,
         ) as span:
-            _push_agent(agent_name)
-            try:
-                if agent_name:
-                    span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
-
-                span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
-                span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
-
-                _set_tools_on_span(span, tools)
-
-                # Run the agent
-                result = f(self, *args, **kwargs)
-
-                input = result.get("input")
-                if (
-                    input is not None
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
-                    normalized_messages = normalize_message_roles([input])
-                    scope = sentry_sdk.get_current_scope()
-                    messages_data = truncate_and_annotate_messages(
-                        normalized_messages, span, scope
+            if agent_name:
+                span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
+
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+            span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+
+            _set_tools_on_span(span, tools)
+
+            # Run the agent
+            result = f(self, *args, **kwargs)
+
+            input = result.get("input")
+            if (
+                input is not None
+                and should_send_default_pii()
+                and integration.include_prompts
+            ):
+                normalized_messages = normalize_message_roles([input])
+                scope = sentry_sdk.get_current_scope()
+                messages_data = truncate_and_annotate_messages(
+                    normalized_messages, span, scope
+                )
+                if messages_data is not None:
+                    set_data_normalized(
+                        span,
+                        SPANDATA.GEN_AI_REQUEST_MESSAGES,
+                        messages_data,
+                        unpack=False,
                     )
-                    if messages_data is not None:
-                        set_data_normalized(
-                            span,
-                            SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                            messages_data,
-                            unpack=False,
-                        )
 
-                output = result.get("output")
-                if (
-                    output is not None
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
-                    set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
+            output = result.get("output")
+            if (
+                output is not None
+                and should_send_default_pii()
+                and integration.include_prompts
+            ):
+                set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
-                return result
-            finally:
-                # Ensure agent is popped even if an exception occurs
-                _pop_agent()
+            return result
 
     return new_invoke
 
@@ -1045,13 +990,11 @@ def new_stream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
         span = start_span_function(
             op=OP.GEN_AI_INVOKE_AGENT,
-            name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
+            name=f"invoke_agent {agent_name}".strip(),
             origin=LangchainIntegration.origin,
         )
         span.__enter__()
 
-        _push_agent(agent_name)
-
         if agent_name:
             span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
 
@@ -1084,57 +1027,43 @@ def new_stream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
 
         old_iterator = result
 
-        def new_iterator() -> "Iterator[Any]":
-            exc_info: "tuple[Any, Any, Any]" = (None, None, None)
-            try:
-                for event in old_iterator:
-                    yield event
+        def new_iterator():
+            # type: () -> Iterator[Any]
+            for event in old_iterator:
+                yield event
 
-                try:
-                    output = event.get("output")
-                except Exception:
-                    output = None
-
-                if (
-                    output is not None
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
-                    set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
-            except Exception:
-                exc_info = sys.exc_info()
-                set_span_errored(span)
-                raise
-            finally:
-                # Ensure cleanup happens even if iterator is abandoned or fails
-                _pop_agent()
-                span.__exit__(*exc_info)
-
-        async def new_iterator_async() -> "AsyncIterator[Any]":
-            exc_info: "tuple[Any, Any, Any]" = (None, None, None)
             try:
-                async for event in old_iterator:
-                    yield event
+                output = event.get("output")
+            except Exception:
+                output = None
 
-                try:
-                    output = event.get("output")
-                except Exception:
-                    output = None
-
-                if (
-                    output is not None
-                    and should_send_default_pii()
-                    and integration.include_prompts
-                ):
-                    set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
+            if (
+                output is not None
+                and should_send_default_pii()
+                and integration.include_prompts
+            ):
+                set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
+
+            span.__exit__(None, None, None)
+
+        async def new_iterator_async():
+            # type: () -> AsyncIterator[Any]
+            async for event in old_iterator:
+                yield event
+
+            try:
+                output = event.get("output")
             except Exception:
-                exc_info = sys.exc_info()
-                set_span_errored(span)
-                raise
-            finally:
-                # Ensure cleanup happens even if iterator is abandoned or fails
-                _pop_agent()
-                span.__exit__(*exc_info)
+                output = None
+
+            if (
+                output is not None
+                and should_send_default_pii()
+                and integration.include_prompts
+            ):
+                set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
+
+            span.__exit__(None, None, None)
 
         if str(type(result)) == "<class 'async_generator'>":
             result = new_iterator_async()

tests/integrations/langchain/test_langchain.py

@@ -66,6 +66,7 @@ def _llm_type(self) -> str:
         return llm_type
 
 
+@pytest.mark.xfail
 @pytest.mark.parametrize(
     "send_default_pii, include_prompts, use_unknown_llm_type",
     [
@@ -222,17 +223,20 @@ def test_langchain_agent(
     # We can't guarantee anything about the "shape" of the langchain execution graph
     assert len(list(x for x in tx["spans"] if x["op"] == "gen_ai.chat")) > 0
 
-    # Token usage is only available in newer versions of langchain (v0.2+)
-    # where usage_metadata is supported on AIMessageChunk
-    if "gen_ai.usage.input_tokens" in chat_spans[0]["data"]:
-        assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142
-        assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50
-        assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192
+    assert "gen_ai.usage.input_tokens" in chat_spans[0]["data"]
+    assert "gen_ai.usage.output_tokens" in chat_spans[0]["data"]
+    assert "gen_ai.usage.total_tokens" in chat_spans[0]["data"]
 
-    if "gen_ai.usage.input_tokens" in chat_spans[1]["data"]:
-        assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
-        assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28
-        assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117
+    assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142
+    assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50
+    assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192
+
+    assert "gen_ai.usage.input_tokens" in chat_spans[1]["data"]
+    assert "gen_ai.usage.output_tokens" in chat_spans[1]["data"]
+    assert "gen_ai.usage.total_tokens" in chat_spans[1]["data"]
+    assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
+    assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28
+    assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117
 
     if send_default_pii and include_prompts:
         assert "5" in chat_spans[0]["data"][SPANDATA.GEN_AI_RESPONSE_TEXT]
@@ -280,8 +284,8 @@ def test_langchain_agent(
         assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_spans[1].get("data", {})
         assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in chat_spans[1].get("data", {})
         assert SPANDATA.GEN_AI_RESPONSE_TEXT not in chat_spans[1].get("data", {})
-        assert SPANDATA.GEN_AI_TOOL_INPUT not in tool_exec_span.get("data", {})
-        assert SPANDATA.GEN_AI_TOOL_OUTPUT not in tool_exec_span.get("data", {})
+        assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in tool_exec_span.get("data", {})
+        assert SPANDATA.GEN_AI_RESPONSE_TEXT not in tool_exec_span.get("data", {})
 
         # Verify tool calls are NOT recorded when PII is disabled
         assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in chat_spans[0].get(