From 5ed981d2e4ee8d3770d9191e3ce8d949b7641635 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 20 Mar 2026 10:49:35 +0100 Subject: [PATCH] revert: fix(integrations): ensure that GEN_AI_AGENT_NAME is properly set for GEN_AI spans under an invoke_agent span --- sentry_sdk/integrations/langchain.py | 211 ++++++------------ .../integrations/langchain/test_langchain.py | 28 ++- 2 files changed, 86 insertions(+), 153 deletions(-) diff --git a/sentry_sdk/integrations/langchain.py b/sentry_sdk/integrations/langchain.py index d19d9bbdd5..ba1794fcdc 100644 --- a/sentry_sdk/integrations/langchain.py +++ b/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)) == "": result = new_iterator_async() diff --git a/tests/integrations/langchain/test_langchain.py b/tests/integrations/langchain/test_langchain.py index 132da0a9a0..6b167c33f1 100644 --- a/tests/integrations/langchain/test_langchain.py +++ b/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(