From 149bd57e24f8a7483a28a8ed9ada197ec6113321 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 4 Sep 2026 15:41:57 +0200 Subject: [PATCH 1/2] chore(langgraph): Remove transaction-based tracing --- sentry_sdk/integrations/langgraph.py | 232 +-- .../integrations/langgraph/test_langgraph.py | 1445 ++++------------- 2 files changed, 395 insertions(+), 1282 deletions(-) diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 3e6c6f7329..172721b39c 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -3,7 +3,6 @@ import sentry_sdk from sentry_sdk.ai.utils import ( - get_start_span_function, normalize_message_roles, set_data_normalized, truncate_and_annotate_messages, @@ -26,7 +25,6 @@ try: from langgraph.errors import GraphBubbleUp - from langgraph.graph import StateGraph from langgraph.pregel import Pregel except ImportError: raise DidNotEnable("langgraph not installed or incompatible") @@ -52,7 +50,6 @@ def setup_once() -> None: # The streaming methods are not patched, because due to some internal reasons, LangGraph # will automatically patch the streaming methods to run through invoke, and by doing this # we prevent duplicate spans for invocations. - StateGraph.compile = _wrap_state_graph_compile(StateGraph.compile) if hasattr(Pregel, "invoke"): Pregel.invoke = _wrap_pregel_invoke(Pregel.invoke) if hasattr(Pregel, "ainvoke"): @@ -138,58 +135,6 @@ def _parse_langgraph_messages(state: "Any") -> "Optional[List[Any]]": return normalized_messages if normalized_messages else None -def _wrap_state_graph_compile(f: "Callable[..., Any]") -> "Callable[..., Any]": - @wraps(f) - def new_compile(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - client = sentry_sdk.get_client() - integration = client.get_integration(LanggraphIntegration) - if integration is None or has_span_streaming_enabled(client.options): - return f(self, *args, **kwargs) - - with sentry_sdk.start_span( - op=OP.GEN_AI_CREATE_AGENT, - origin=LanggraphIntegration.origin, - ) as span: - compiled_graph = f(self, *args, **kwargs) - - compiled_graph_name = getattr(compiled_graph, "name", None) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "create_agent") - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, compiled_graph_name) - - if compiled_graph_name: - span.description = f"create_agent {compiled_graph_name}" - else: - span.description = "create_agent" - - if kwargs.get("model", None) is not None: - span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, kwargs.get("model")) - - tools = None - get_graph = getattr(compiled_graph, "get_graph", None) - if get_graph and callable(get_graph): - graph_obj = compiled_graph.get_graph() - nodes = getattr(graph_obj, "nodes", None) - if nodes and isinstance(nodes, dict): - tools_node = nodes.get("tools") - if tools_node: - data = getattr(tools_node, "data", None) - if data and hasattr(data, "tools_by_name"): - tools = list(data.tools_by_name.keys()) - - if tools is not None: - # Available tools aren't gated on the legacy PII settings, so they're - # only gated when data collection has been configured. - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["inputs"]: - span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools) - else: - span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools) - - return compiled_graph - - return new_compile - - def _wrap_pregel_invoke(f: "Callable[..., Any]") -> "Callable[..., Any]": @wraps(f) def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": @@ -199,95 +144,47 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": return f(self, *args, **kwargs) graph_name = _get_graph_name(self) - span_name = ( - f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" - ) - if has_span_streaming_enabled(client.options): - with sentry_sdk.traces.start_span( - name=span_name, - attributes={ - "sentry.op": OP.GEN_AI_INVOKE_AGENT, - "sentry.origin": LanggraphIntegration.origin, - SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", - }, - ) as span: - if graph_name: - span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) - span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name) - - # Store input messages to later compare with output - input_messages = None - if len(args) > 0: - input_messages = _parse_langgraph_messages(args[0]) - if input_messages and _should_record_inputs(integration): - normalized_input_messages = normalize_message_roles( - input_messages - ) - - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages( - normalized_input_messages, span, scope - ) - if not has_span_streaming_enabled(client.options) - else normalized_input_messages - ) - if messages_data is not None: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - messages_data, - unpack=False, - ) - - result = f(self, *args, **kwargs) + with sentry_sdk.traces.start_span( + name=f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent", + attributes={ + "sentry.op": OP.GEN_AI_INVOKE_AGENT, + "sentry.origin": LanggraphIntegration.origin, + SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", + }, + ) as span: + if graph_name: + span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name) - _set_response_attributes(span, input_messages, result, integration) + # Store input messages to later compare with output + input_messages = None + if len(args) > 0: + input_messages = _parse_langgraph_messages(args[0]) + if input_messages and _should_record_inputs(integration): + normalized_input_messages = normalize_message_roles(input_messages) - return result - else: - with get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name=span_name, - origin=LanggraphIntegration.origin, - ) as span: - if graph_name: - span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") - - # Store input messages to later compare with output - input_messages = None - if len(args) > 0: - input_messages = _parse_langgraph_messages(args[0]) - if input_messages and _should_record_inputs(integration): - normalized_input_messages = normalize_message_roles( - input_messages + scope = sentry_sdk.get_current_scope() + messages_data = ( + truncate_and_annotate_messages( + normalized_input_messages, span, scope ) - - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages( - normalized_input_messages, span, scope - ) - if not has_span_streaming_enabled(client.options) - else normalized_input_messages + if not has_span_streaming_enabled(client.options) + else normalized_input_messages + ) + 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, - ) - result = f(self, *args, **kwargs) + result = f(self, *args, **kwargs) - _set_response_attributes(span, input_messages, result, integration) + _set_response_attributes(span, input_messages, result, integration) - return result + return result return new_invoke @@ -301,63 +198,18 @@ async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": return await f(self, *args, **kwargs) graph_name = _get_graph_name(self) - span_name = ( - f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" - ) - if has_span_streaming_enabled(client.options): - with sentry_sdk.traces.start_span( - name=span_name, - attributes={ - "sentry.op": OP.GEN_AI_INVOKE_AGENT, - "sentry.origin": LanggraphIntegration.origin, - SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", - }, - ) as span: - if graph_name: - span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) - span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name) - - input_messages = None - if len(args) > 0: - input_messages = _parse_langgraph_messages(args[0]) - if input_messages and _should_record_inputs(integration): - normalized_input_messages = normalize_message_roles( - input_messages - ) - - scope = sentry_sdk.get_current_scope() - messages_data = ( - truncate_and_annotate_messages( - normalized_input_messages, span, scope - ) - if not has_span_streaming_enabled(client.options) - else normalized_input_messages - ) - if messages_data is not None: - set_data_normalized( - span, - SPANDATA.GEN_AI_REQUEST_MESSAGES, - messages_data, - unpack=False, - ) - - result = await f(self, *args, **kwargs) - - _set_response_attributes(span, input_messages, result, integration) - - return result - - with get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name=span_name, - origin=LanggraphIntegration.origin, + with sentry_sdk.traces.start_span( + name=f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent", + attributes={ + "sentry.op": OP.GEN_AI_INVOKE_AGENT, + "sentry.origin": LanggraphIntegration.origin, + SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", + }, ) as span: if graph_name: - span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name) - - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name) input_messages = None if len(args) > 0: diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index 1a9d3919b4..76055f9379 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -1,7 +1,7 @@ import asyncio import json import sys -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest from langchain_core.language_models.chat_models import BaseChatModel @@ -39,7 +39,6 @@ def mock_langgraph_imports(): _parse_langgraph_messages, _wrap_pregel_ainvoke, _wrap_pregel_invoke, - _wrap_state_graph_compile, ) @@ -152,66 +151,6 @@ def test_langgraph_integration_init(): assert integration.origin == "auto.ai.langgraph" -@pytest.mark.parametrize( - "send_default_pii, include_prompts", - [ - (True, True), - (True, False), - (False, True), - (False, False), - ], -) -def test_state_graph_compile( - sentry_init, - capture_events, - send_default_pii, - include_prompts, -): - """Test StateGraph.compile() wrapper creates proper create_agent span.""" - sentry_init( - integrations=[LanggraphIntegration(include_prompts=include_prompts)], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - stream_gen_ai_spans=False, - ) - - graph = MockStateGraph() - - def original_compile(self, *args, **kwargs): - return MockCompiledGraph(self.name) - - events = capture_events() - - with patch("sentry_sdk.integrations.langgraph.StateGraph"), start_transaction(): - wrapped_compile = _wrap_state_graph_compile(original_compile) - compiled_graph = wrapped_compile(graph, model="test-model", checkpointer=None) - - assert compiled_graph is not None - assert compiled_graph.name == "test_graph" - - tx = events[0] - assert tx["type"] == "transaction" - - agent_spans = [span for span in tx["spans"] if span["op"] == OP.GEN_AI_CREATE_AGENT] - assert len(agent_spans) == 1 - agent_span = agent_spans[0] - - assert agent_span["description"] == "create_agent test_graph" - assert agent_span["origin"] == "auto.ai.langgraph" - assert agent_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "create_agent" - assert agent_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" - assert agent_span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "test-model" - assert SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in agent_span["data"] - - tools_data = agent_span["data"][SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] - - assert tools_data == ["search_tool", "calculator"] - assert len(tools_data) == 2 - assert "search_tool" in tools_data - assert "calculator" in tools_data - - -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -223,19 +162,16 @@ def original_compile(self, *args, **kwargs): ) def test_pregel_invoke( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, - span_streaming, ): """Test Pregel.invoke() wrapper creates proper invoke_agent span.""" sentry_init( integrations=[LanggraphIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -267,127 +203,62 @@ def original_invoke(self, *args, **kwargs): ] return {"messages": new_messages} - if span_streaming: - items = capture_items("span") + items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + assert result is not None - invoke_span = invoke_spans[0] - assert invoke_span["name"] == "invoke_agent test_graph" - assert invoke_span["attributes"]["sentry.origin"] == "auto.ai.langgraph" - assert ( - invoke_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" - ) - assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "test_graph" - assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - if send_default_pii and include_prompts: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] + invoke_span = invoke_spans[0] + assert invoke_span["name"] == "invoke_agent test_graph" + assert invoke_span["attributes"]["sentry.origin"] == "auto.ai.langgraph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "test_graph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" + + if send_default_pii and include_prompts: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] - request_messages = invoke_span["attributes"][ - SPANDATA.GEN_AI_REQUEST_MESSAGES - ] + request_messages = invoke_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - if isinstance(request_messages, str): - request_messages = json.loads(request_messages) - assert len(request_messages) == 2 - assert request_messages[0]["content"] == "Hello, can you help me?" - assert request_messages[1]["content"] == "Of course! How can I assist you?" + if isinstance(request_messages, str): + request_messages = json.loads(request_messages) + assert len(request_messages) == 2 + assert request_messages[0]["content"] == "Hello, can you help me?" + assert request_messages[1]["content"] == "Of course! How can I assist you?" - response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == expected_assistant_response + response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == expected_assistant_response - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] - tool_calls_data = invoke_span["attributes"][ - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS - ] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] + tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - if isinstance(tool_calls_data, str): - tool_calls_data = json.loads(tool_calls_data) + if isinstance(tool_calls_data, str): + tool_calls_data = json.loads(tool_calls_data) - assert len(tool_calls_data) == 1 - assert tool_calls_data[0]["id"] == "call_test_123" - assert tool_calls_data[0]["function"]["name"] == "search_tool" - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get( - "attributes", {} - ) - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get( - "attributes", {} - ) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( - "attributes", {} - ) + assert len(tool_calls_data) == 1 + assert tool_calls_data[0]["id"] == "call_test_123" + assert tool_calls_data[0]["function"]["name"] == "search_tool" else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - tx = events[0] - assert tx["type"] == "transaction" - - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_span = invoke_spans[0] - assert invoke_span["description"] == "invoke_agent test_graph" - assert invoke_span["origin"] == "auto.ai.langgraph" - assert invoke_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" - assert invoke_span["data"][SPANDATA.GEN_AI_PIPELINE_NAME] == "test_graph" - assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" - - if send_default_pii and include_prompts: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["data"] - - request_messages = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - - if isinstance(request_messages, str): - request_messages = json.loads(request_messages) - assert len(request_messages) == 1 - assert request_messages[0]["content"] == "Of course! How can I assist you?" - - response_text = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == expected_assistant_response - - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["data"] - tool_calls_data = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - - if isinstance(tool_calls_data, str): - tool_calls_data = json.loads(tool_calls_data) - - assert len(tool_calls_data) == 1 - assert tool_calls_data[0]["id"] == "call_test_123" - assert tool_calls_data[0]["function"]["name"] == "search_tool" - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("data", {}) - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("data", {}) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( - "data", {} - ) + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( + "attributes", {} + ) -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -399,19 +270,16 @@ def original_invoke(self, *args, **kwargs): ) def test_pregel_ainvoke( sentry_init, - capture_events, capture_items, send_default_pii, include_prompts, - span_streaming, ): """Test Pregel.ainvoke() async wrapper creates proper invoke_agent span.""" sentry_init( integrations=[LanggraphIntegration(include_prompts=include_prompts)], traces_sample_rate=1.0, send_default_pii=send_default_pii, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = {"messages": [MockMessage("What's the weather like?", name="user")]} @@ -443,117 +311,61 @@ async def run_test(): result = await wrapped_ainvoke(pregel, test_state) return result - if span_streaming: - items = capture_items("span") + items = capture_items("span") - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + result = asyncio.run(run_test()) + assert result is not None - invoke_span = invoke_spans[0] - assert invoke_span["name"] == "invoke_agent async_graph" - assert invoke_span["attributes"]["sentry.origin"] == "auto.ai.langgraph" - assert ( - invoke_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" - ) - assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "async_graph" - assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "async_graph" + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - if send_default_pii and include_prompts: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] + invoke_span = invoke_spans[0] + assert invoke_span["name"] == "invoke_agent async_graph" + assert invoke_span["attributes"]["sentry.origin"] == "auto.ai.langgraph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == "async_graph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == "async_graph" + + if send_default_pii and include_prompts: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] - response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == expected_assistant_response + response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == expected_assistant_response - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] - tool_calls_data = invoke_span["attributes"][ - SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS - ] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] + tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - if isinstance(tool_calls_data, str): - tool_calls_data = json.loads(tool_calls_data) + if isinstance(tool_calls_data, str): + tool_calls_data = json.loads(tool_calls_data) - assert len(tool_calls_data) == 1 - assert tool_calls_data[0]["id"] == "call_weather_456" - assert tool_calls_data[0]["function"]["name"] == "get_weather" - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get( - "attributes", {} - ) - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get( - "attributes", {} - ) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( - "attributes", {} - ) + assert len(tool_calls_data) == 1 + assert tool_calls_data[0]["id"] == "call_weather_456" + assert tool_calls_data[0]["function"]["name"] == "get_weather" else: - events = capture_events() - - result = asyncio.run(run_test()) - assert result is not None - - tx = events[0] - assert tx["type"] == "transaction" - - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_span = invoke_spans[0] - assert invoke_span["description"] == "invoke_agent async_graph" - assert invoke_span["origin"] == "auto.ai.langgraph" - assert invoke_span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" - assert invoke_span["data"][SPANDATA.GEN_AI_PIPELINE_NAME] == "async_graph" - assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == "async_graph" - - if send_default_pii and include_prompts: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["data"] - - response_text = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == expected_assistant_response - - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["data"] - tool_calls_data = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - - if isinstance(tool_calls_data, str): - tool_calls_data = json.loads(tool_calls_data) - - assert len(tool_calls_data) == 1 - assert tool_calls_data[0]["id"] == "call_weather_456" - assert tool_calls_data[0]["function"]["name"] == "get_weather" - else: - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("data", {}) - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("data", {}) - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( - "data", {} - ) + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in invoke_span.get( + "attributes", {} + ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_invoke_error( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test error handling during graph execution.""" sentry_init( integrations=[LanggraphIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = {"messages": [MockMessage("This will fail")]} @@ -562,60 +374,35 @@ def test_pregel_invoke_error( def original_invoke(self, *args, **kwargs): raise Exception("Graph execution failed") - if span_streaming: - items = capture_items("span") - - with start_transaction(), pytest.raises( - Exception, match="Graph execution failed" - ): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + items = capture_items("span") - invoke_span = invoke_spans[0] - assert invoke_span.get("status") == "error" - else: - events = capture_events() - - with start_transaction(), pytest.raises( - Exception, match="Graph execution failed" - ): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) + with start_transaction(), pytest.raises(Exception, match="Graph execution failed"): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) - tx = events[0] - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_span = invoke_spans[0] - assert invoke_span.get("status") == "internal_error" - assert invoke_span.get("tags", {}).get("status") == "internal_error" + invoke_span = invoke_spans[0] + assert invoke_span.get("status") == "error" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_ainvoke_error( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test error handling during async graph execution.""" sentry_init( integrations=[LanggraphIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = {"messages": [MockMessage("This will fail async")]} @@ -631,85 +418,35 @@ async def run_error_test(): wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) await wrapped_ainvoke(pregel, test_state) - if span_streaming: - items = capture_items("span") - - asyncio.run(run_error_test()) - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_span = invoke_spans[0] - assert invoke_span.get("status") == "error" - else: - events = capture_events() - - asyncio.run(run_error_test()) - - tx = events[0] - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_span = invoke_spans[0] - assert invoke_span.get("status") == "internal_error" - assert invoke_span.get("tags", {}).get("status") == "internal_error" - - -def test_span_origin( - sentry_init, - capture_events, -): - """Test that span origins are correctly set.""" - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - stream_gen_ai_spans=False, - ) - - graph = MockStateGraph() - - def original_compile(self, *args, **kwargs): - return MockCompiledGraph(self.name) - - events = capture_events() - - with start_transaction(): - from sentry_sdk.integrations.langgraph import _wrap_state_graph_compile + items = capture_items("span") - wrapped_compile = _wrap_state_graph_compile(original_compile) - wrapped_compile(graph) + asyncio.run(run_error_test()) - tx = events[0] - assert tx["contexts"]["trace"]["origin"] == "manual" + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - for span in tx["spans"]: - assert span["origin"] == "auto.ai.langgraph" + invoke_span = invoke_spans[0] + assert invoke_span.get("status") == "error" -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("graph_name", ["my_graph", None, ""]) def test_pregel_invoke_with_different_graph_names( sentry_init, - capture_events, capture_items, graph_name, - span_streaming, ): """Test Pregel.invoke() with different graph name scenarios.""" sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) pregel = MockPregelInstance(graph_name) if graph_name else MockPregelInstance() @@ -720,67 +457,36 @@ def test_pregel_invoke_with_different_graph_names( def original_invoke(self, *args, **kwargs): return {"result": "test"} - if span_streaming: - items = capture_items("span") - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, {"messages": []}) - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_span = invoke_spans[0] - - if graph_name and graph_name.strip(): - assert invoke_span["name"] == "invoke_agent my_graph" - assert ( - invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == graph_name - ) - assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == graph_name - else: - assert invoke_span["name"] == "invoke_agent" - assert SPANDATA.GEN_AI_PIPELINE_NAME not in invoke_span.get( - "attributes", {} - ) - assert SPANDATA.GEN_AI_AGENT_NAME not in invoke_span.get("attributes", {}) - else: - events = capture_events() + items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, {"messages": []}) + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, {"messages": []}) - tx = events[0] - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_span = invoke_spans[0] + invoke_span = invoke_spans[0] - if graph_name and graph_name.strip(): - assert invoke_span["description"] == "invoke_agent my_graph" - assert invoke_span["data"][SPANDATA.GEN_AI_PIPELINE_NAME] == graph_name - assert invoke_span["data"][SPANDATA.GEN_AI_AGENT_NAME] == graph_name - else: - assert invoke_span["description"] == "invoke_agent" - assert SPANDATA.GEN_AI_PIPELINE_NAME not in invoke_span.get("data", {}) - assert SPANDATA.GEN_AI_AGENT_NAME not in invoke_span.get("data", {}) + if graph_name and graph_name.strip(): + assert invoke_span["name"] == "invoke_agent my_graph" + assert invoke_span["attributes"][SPANDATA.GEN_AI_PIPELINE_NAME] == graph_name + assert invoke_span["attributes"][SPANDATA.GEN_AI_AGENT_NAME] == graph_name + else: + assert invoke_span["name"] == "invoke_agent" + assert SPANDATA.GEN_AI_PIPELINE_NAME not in invoke_span.get("attributes", {}) + assert SPANDATA.GEN_AI_AGENT_NAME not in invoke_span.get("attributes", {}) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_invoke_span_includes_usage_data( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent spans include aggregated usage data from context_wrapper. @@ -789,8 +495,7 @@ def test_pregel_invoke_span_includes_usage_data( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -830,73 +535,40 @@ def original_invoke(self, *args, **kwargs): ] return {"messages": new_messages} - if span_streaming: - items = capture_items("span") - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has usage data - assert invoke_agent_span["name"] == "invoke_agent test_graph" - assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"] - - # The usage should match the mock_usage values (aggregated across all calls) - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30 - else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + items = capture_items("span") - assert result is not None + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - tx = events[0] - assert tx["type"] == "transaction" + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span has usage data - assert invoke_agent_span["description"] == "invoke_agent test_graph" - assert "gen_ai.usage.input_tokens" in invoke_agent_span["data"] - assert "gen_ai.usage.output_tokens" in invoke_agent_span["data"] - assert "gen_ai.usage.total_tokens" in invoke_agent_span["data"] + # Verify invoke_agent span has usage data + assert invoke_agent_span["name"] == "invoke_agent test_graph" + assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"] + assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"] + assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"] - # The usage should match the mock_usage values (aggregated across all calls) - assert invoke_agent_span["data"]["gen_ai.usage.input_tokens"] == 10 - assert invoke_agent_span["data"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["data"]["gen_ai.usage.total_tokens"] == 30 + # The usage should match the mock_usage values (aggregated across all calls) + assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10 + assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 + assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_ainvoke_span_includes_usage_data( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent spans include aggregated usage data from context_wrapper. @@ -905,8 +577,7 @@ def test_pregel_ainvoke_span_includes_usage_data( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -952,67 +623,37 @@ async def run_test(): result = await wrapped_ainvoke(pregel, test_state) return result - if span_streaming: - items = capture_items("span") - - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has usage data - assert invoke_agent_span["name"] == "invoke_agent test_graph" - assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"] - - # The usage should match the mock_usage values (aggregated across all calls) - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30 - else: - events = capture_events() + items = capture_items("span") - result = asyncio.run(run_test()) - assert result is not None - - tx = events[0] - assert tx["type"] == "transaction" + result = asyncio.run(run_test()) + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span has usage data - assert invoke_agent_span["description"] == "invoke_agent test_graph" - assert "gen_ai.usage.input_tokens" in invoke_agent_span["data"] - assert "gen_ai.usage.output_tokens" in invoke_agent_span["data"] - assert "gen_ai.usage.total_tokens" in invoke_agent_span["data"] + # Verify invoke_agent span has usage data + assert invoke_agent_span["name"] == "invoke_agent test_graph" + assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"] + assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"] + assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"] - # The usage should match the mock_usage values (aggregated across all calls) - assert invoke_agent_span["data"]["gen_ai.usage.input_tokens"] == 10 - assert invoke_agent_span["data"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["data"]["gen_ai.usage.total_tokens"] == 30 + # The usage should match the mock_usage values (aggregated across all calls) + assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10 + assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 + assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_invoke_multiple_llm_calls_aggregate_usage( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent spans show aggregated usage across multiple LLM calls @@ -1021,8 +662,7 @@ def test_pregel_invoke_multiple_llm_calls_aggregate_usage( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -1073,61 +713,34 @@ def original_invoke(self, *args, **kwargs): ] return {"messages": new_messages} - if span_streaming: - items = capture_items("span") - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None + items = capture_items("span") - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has aggregated usage from both API calls - # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50 - else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - tx = events[0] - assert tx["type"] == "transaction" + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span has aggregated usage from both API calls - # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total - assert invoke_agent_span["data"]["gen_ai.usage.input_tokens"] == 30 - assert invoke_agent_span["data"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["data"]["gen_ai.usage.total_tokens"] == 50 + # Verify invoke_agent span has aggregated usage from both API calls + # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total + assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30 + assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 + assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_ainvoke_multiple_llm_calls_aggregate_usage( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent spans show aggregated usage across multiple LLM calls @@ -1136,8 +749,7 @@ def test_pregel_ainvoke_multiple_llm_calls_aggregate_usage( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -1194,55 +806,31 @@ async def run_test(): result = await wrapped_ainvoke(pregel, test_state) return result - if span_streaming: - items = capture_items("span") + items = capture_items("span") - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has aggregated usage from both API calls - # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50 - else: - events = capture_events() - - result = asyncio.run(run_test()) - assert result is not None - - tx = events[0] - assert tx["type"] == "transaction" + result = asyncio.run(run_test()) + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span has aggregated usage from both API calls - # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total - assert invoke_agent_span["data"]["gen_ai.usage.input_tokens"] == 30 - assert invoke_agent_span["data"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["data"]["gen_ai.usage.total_tokens"] == 50 + # Verify invoke_agent span has aggregated usage from both API calls + # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total + assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30 + assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 + assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50 -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_invoke_span_includes_response_model( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent spans include the response model. @@ -1251,8 +839,7 @@ def test_pregel_invoke_span_includes_response_model( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -1292,66 +879,36 @@ def original_invoke(self, *args, **kwargs): ] return {"messages": new_messages} - if span_streaming: - items = capture_items("span") + items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has response model - assert invoke_agent_span["name"] == "invoke_agent test_graph" - assert "gen_ai.response.model" in invoke_agent_span["attributes"] - assert ( - invoke_agent_span["attributes"]["gen_ai.response.model"] - == "gpt-4.1-2025-04-14" - ) - else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - tx = events[0] - assert tx["type"] == "transaction" + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span has response model - assert invoke_agent_span["description"] == "invoke_agent test_graph" - assert "gen_ai.response.model" in invoke_agent_span["data"] - assert ( - invoke_agent_span["data"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" - ) + # Verify invoke_agent span has response model + assert invoke_agent_span["name"] == "invoke_agent test_graph" + assert "gen_ai.response.model" in invoke_agent_span["attributes"] + assert ( + invoke_agent_span["attributes"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" + ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_ainvoke_span_includes_response_model( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that invoke_agent spans include the response model. @@ -1360,8 +917,7 @@ def test_pregel_ainvoke_span_includes_response_model( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -1407,60 +963,33 @@ async def run_test(): result = await wrapped_ainvoke(pregel, test_state) return result - if span_streaming: - items = capture_items("span") - - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has response model - assert invoke_agent_span["name"] == "invoke_agent test_graph" - assert "gen_ai.response.model" in invoke_agent_span["attributes"] - assert ( - invoke_agent_span["attributes"]["gen_ai.response.model"] - == "gpt-4.1-2025-04-14" - ) - else: - events = capture_events() - - result = asyncio.run(run_test()) - assert result is not None + items = capture_items("span") - tx = events[0] - assert tx["type"] == "transaction" + result = asyncio.run(run_test()) + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span has response model - assert invoke_agent_span["description"] == "invoke_agent test_graph" - assert "gen_ai.response.model" in invoke_agent_span["data"] - assert ( - invoke_agent_span["data"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" - ) + # Verify invoke_agent span has response model + assert invoke_agent_span["name"] == "invoke_agent test_graph" + assert "gen_ai.response.model" in invoke_agent_span["attributes"] + assert ( + invoke_agent_span["attributes"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" + ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_invoke_span_uses_last_response_model( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that when an agent makes multiple LLM calls (e.g., with tools), @@ -1469,8 +998,7 @@ def test_pregel_invoke_span_uses_last_response_model( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -1523,64 +1051,35 @@ def original_invoke(self, *args, **kwargs): ] return {"messages": new_messages} - if span_streaming: - items = capture_items("span") - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None + items = capture_items("span") - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span uses the LAST response model - assert "gen_ai.response.model" in invoke_agent_span["attributes"] - assert ( - invoke_agent_span["attributes"]["gen_ai.response.model"] - == "gpt-4.1-2025-04-14" - ) - else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - tx = events[0] - assert tx["type"] == "transaction" + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span uses the LAST response model - assert "gen_ai.response.model" in invoke_agent_span["data"] - assert ( - invoke_agent_span["data"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" - ) + # Verify invoke_agent span uses the LAST response model + assert "gen_ai.response.model" in invoke_agent_span["attributes"] + assert ( + invoke_agent_span["attributes"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" + ) -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_ainvoke_span_uses_last_response_model( sentry_init, - capture_events, capture_items, - span_streaming, ): """ Test that when an agent makes multiple LLM calls (e.g., with tools), @@ -1589,8 +1088,7 @@ def test_pregel_ainvoke_span_uses_last_response_model( sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) test_state = { @@ -1649,50 +1147,27 @@ async def run_test(): result = await wrapped_ainvoke(pregel, test_state) return result - if span_streaming: - items = capture_items("span") - - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span uses the LAST response model - assert "gen_ai.response.model" in invoke_agent_span["attributes"] - assert ( - invoke_agent_span["attributes"]["gen_ai.response.model"] - == "gpt-4.1-2025-04-14" - ) - else: - events = capture_events() - - result = asyncio.run(run_test()) - assert result is not None + items = capture_items("span") - tx = events[0] - assert tx["type"] == "transaction" + result = asyncio.run(run_test()) + assert result is not None - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] + invoke_agent_span = invoke_spans[0] - # Verify invoke_agent span uses the LAST response model - assert "gen_ai.response.model" in invoke_agent_span["data"] - assert ( - invoke_agent_span["data"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" - ) + # Verify invoke_agent span uses the LAST response model + assert "gen_ai.response.model" in invoke_agent_span["attributes"] + assert ( + invoke_agent_span["attributes"]["gen_ai.response.model"] == "gpt-4.1-2025-04-14" + ) def test_complex_message_parsing(): @@ -1742,20 +1217,16 @@ def test_complex_message_parsing(): assert result[2]["function_call"]["name"] == "search" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_extraction_functions_complex_scenario( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test extraction functions with complex scenarios including multiple messages and edge cases.""" sentry_init( integrations=[LanggraphIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) pregel = MockPregelInstance("complex_graph") @@ -1791,55 +1262,31 @@ def original_invoke(self, *args, **kwargs): ] return {"messages": new_messages} - if span_streaming: - items = capture_items("span") - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + items = capture_items("span") - invoke_span = invoke_spans[0] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] - response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == "Final response" - - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] - - tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] - else: - events = capture_events() - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + with start_transaction(): + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) - assert result is not None + assert result is not None - tx = events[0] - invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 + sentry_sdk.flush() + spans = [item.payload for item in items] + invoke_spans = [ + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT + ] + assert len(invoke_spans) == 1 - invoke_span = invoke_spans[0] - assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["data"] - response_text = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert response_text == "Final response" + invoke_span = invoke_spans[0] + assert SPANDATA.GEN_AI_RESPONSE_TEXT in invoke_span["attributes"] + response_text = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert response_text == "Final response" - assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS in invoke_span["attributes"] - tool_calls_data = invoke_span["data"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + tool_calls_data = invoke_span["attributes"][SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] if isinstance(tool_calls_data, str): tool_calls_data = json.loads(tool_calls_data) @@ -1851,20 +1298,16 @@ def original_invoke(self, *args, **kwargs): assert tool_calls_data[1]["function"]["name"] == "calculate" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_langgraph_message_role_mapping( sentry_init, - capture_events, capture_items, - span_streaming, ): """Test that Langgraph integration properly maps message roles like 'ai' to 'assistant'""" sentry_init( integrations=[LanggraphIntegration(include_prompts=True)], traces_sample_rate=1.0, send_default_pii=True, - trace_lifecycle="stream" if span_streaming else "static", - stream_gen_ai_spans=False, + trace_lifecycle="stream", ) # Mock a langgraph message with mixed roles @@ -1885,51 +1328,28 @@ def __init__(self, content, message_type="human"): compiled_graph = MockCompiledGraph("test_graph") pregel = MockPregelInstance(compiled_graph) + items = capture_items("span") - if span_streaming: - items = capture_items("span") - - with start_transaction(name="langgraph tx"): - # Use the wrapped invoke function directly - from sentry_sdk.integrations.langgraph import _wrap_pregel_invoke - - wrapped_invoke = _wrap_pregel_invoke( - lambda self, state_data: {"result": "success"} - ) - wrapped_invoke(pregel, state_data) - - sentry_sdk.flush() - span = next(item.payload for item in items) - - # Verify that the span was created correctly - assert span["attributes"]["sentry.op"] == "gen_ai.invoke_agent" + with start_transaction(name="langgraph tx"): + # Use the wrapped invoke function directly + from sentry_sdk.integrations.langgraph import _wrap_pregel_invoke - # If messages were captured, verify role mapping - if SPANDATA.GEN_AI_REQUEST_MESSAGES in span["attributes"]: - stored_messages = json.loads( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - else: - events = capture_events() - - with start_transaction(name="langgraph tx"): - # Use the wrapped invoke function directly - from sentry_sdk.integrations.langgraph import _wrap_pregel_invoke - - wrapped_invoke = _wrap_pregel_invoke( - lambda self, state_data: {"result": "success"} - ) - wrapped_invoke(pregel, state_data) + wrapped_invoke = _wrap_pregel_invoke( + lambda self, state_data: {"result": "success"} + ) + wrapped_invoke(pregel, state_data) - (event,) = events - span = event["spans"][0] + sentry_sdk.flush() + span = next(item.payload for item in items) - # Verify that the span was created correctly - assert span["op"] == "gen_ai.invoke_agent" + # Verify that the span was created correctly + assert span["attributes"]["sentry.op"] == "gen_ai.invoke_agent" - # If messages were captured, verify role mapping - if SPANDATA.GEN_AI_REQUEST_MESSAGES in span["data"]: - stored_messages = json.loads(span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES]) + # If messages were captured, verify role mapping + if SPANDATA.GEN_AI_REQUEST_MESSAGES in span["attributes"]: + stored_messages = json.loads( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) # Find messages with specific content to verify role mapping ai_message = next( @@ -1953,66 +1373,9 @@ def __init__(self, content, message_type="human"): assert "ai" not in roles -def test_langgraph_message_truncation(sentry_init, capture_events): - """Test that large messages are truncated properly in Langgraph integration.""" - - sentry_init( - integrations=[LanggraphIntegration(include_prompts=True)], - traces_sample_rate=1.0, - send_default_pii=True, - stream_gen_ai_spans=False, - ) - events = capture_events() - - large_content = ( - "This is a very long message that will exceed our size limits. " * 1000 - ) - test_state = { - "messages": [ - MockMessage("small message 1", name="user"), - MockMessage(large_content, name="assistant"), - MockMessage(large_content, name="user"), - MockMessage("small message 4", name="assistant"), - MockMessage("small message 5", name="user"), - ] - } - - pregel = MockPregelInstance("test_graph") - - def original_invoke(self, *args, **kwargs): - return {"messages": args[0].get("messages", [])} - - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - assert len(events) > 0 - tx = events[0] - assert tx["type"] == "transaction" - - invoke_spans = [ - span for span in tx.get("spans", []) if span.get("op") == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) > 0 - - invoke_span = invoke_spans[0] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES in invoke_span["data"] - - messages_data = invoke_span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert isinstance(messages_data, str) - - parsed_messages = json.loads(messages_data) - assert isinstance(parsed_messages, list) - assert len(parsed_messages) == 1 - assert "small message 5" in str(parsed_messages[0]) - assert tx["_meta"]["spans"]["0"]["data"]["gen_ai.request.messages"][""]["len"] == 5 - - def test_graph_bubble_up_ignored(sentry_init, capture_items): sentry_init( integrations=[LanggraphIntegration()], - stream_gen_ai_spans=False, ) events = capture_items("event") @@ -2024,27 +1387,18 @@ def test_graph_bubble_up_ignored(sentry_init, capture_items): assert len(events) == 0 -def _invoke_span_data(items_or_events, span_streaming): - if span_streaming: - sentry_sdk.flush() - spans = [item.payload for item in items_or_events] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - return invoke_spans[0]["attributes"] - - tx = items_or_events[0] +def _invoke_span_data(items_or_events): + sentry_sdk.flush() + spans = [item.payload for item in items_or_events] invoke_spans = [ - span for span in tx["spans"] if span["op"] == OP.GEN_AI_INVOKE_AGENT + span + for span in spans + if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT ] assert len(invoke_spans) == 1 - return invoke_spans[0]["data"] + return invoke_spans[0]["attributes"] -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection, send_default_pii, expect_inputs", [ @@ -2088,19 +1442,16 @@ def _invoke_span_data(items_or_events, span_streaming): ) def test_pregel_invoke_gates_request_messages_on_inputs_setting( sentry_init, - capture_events, capture_items, data_collection, send_default_pii, expect_inputs, - span_streaming, ): init_kwargs = { "integrations": [LanggraphIntegration()], "traces_sample_rate": 1.0, "send_default_pii": send_default_pii, - "trace_lifecycle": "stream" if span_streaming else "static", - "stream_gen_ai_spans": False, + "trace_lifecycle": "stream", } if data_collection is not None: init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -2121,13 +1472,13 @@ def original_invoke(self, *args, **kwargs): ] } - captured = capture_items("span") if span_streaming else capture_events() + captured = capture_items("span") with start_transaction(): wrapped_invoke = _wrap_pregel_invoke(original_invoke) wrapped_invoke(pregel, test_state) - data = _invoke_span_data(captured, span_streaming) + data = _invoke_span_data(captured) if expect_inputs: assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data @@ -2135,7 +1486,6 @@ def original_invoke(self, *args, **kwargs): assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection, send_default_pii, expect_outputs", [ @@ -2179,19 +1529,16 @@ def original_invoke(self, *args, **kwargs): ) def test_pregel_invoke_gates_response_text_and_tool_calls_on_outputs_setting( sentry_init, - capture_events, capture_items, data_collection, send_default_pii, expect_outputs, - span_streaming, ): init_kwargs = { "integrations": [LanggraphIntegration()], "traces_sample_rate": 1.0, "send_default_pii": send_default_pii, - "trace_lifecycle": "stream" if span_streaming else "static", - "stream_gen_ai_spans": False, + "trace_lifecycle": "stream", } if data_collection is not None: init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -2221,13 +1568,13 @@ def original_invoke(self, *args, **kwargs): ] } - captured = capture_items("span") if span_streaming else capture_events() + captured = capture_items("span") with start_transaction(): wrapped_invoke = _wrap_pregel_invoke(original_invoke) wrapped_invoke(pregel, test_state) - data = _invoke_span_data(captured, span_streaming) + data = _invoke_span_data(captured) if expect_outputs: assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response @@ -2239,7 +1586,6 @@ def original_invoke(self, *args, **kwargs): assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data -@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize( "data_collection, send_default_pii, expect_inputs, expect_outputs", [ @@ -2275,20 +1621,17 @@ def original_invoke(self, *args, **kwargs): ) def test_pregel_ainvoke_gates_inputs_and_outputs_independently( sentry_init, - capture_events, capture_items, data_collection, send_default_pii, expect_inputs, expect_outputs, - span_streaming, ): init_kwargs = { "integrations": [LanggraphIntegration()], "traces_sample_rate": 1.0, "send_default_pii": send_default_pii, - "trace_lifecycle": "stream" if span_streaming else "static", - "stream_gen_ai_spans": False, + "trace_lifecycle": "stream", } if data_collection is not None: init_kwargs["_experiments"] = {"data_collection": data_collection} @@ -2323,11 +1666,11 @@ async def run_test(): wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) return await wrapped_ainvoke(pregel, test_state) - captured = capture_items("span") if span_streaming else capture_events() + captured = capture_items("span") asyncio.run(run_test()) - data = _invoke_span_data(captured, span_streaming) + data = _invoke_span_data(captured) if expect_inputs: assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data @@ -2344,21 +1687,17 @@ async def run_test(): assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_invoke_message_delta_ignores_gen_ai_inputs_setting( sentry_init, - capture_events, capture_items, - span_streaming, ): sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", _experiments={ "data_collection": {"gen_ai": {"inputs": False, "outputs": True}} }, - stream_gen_ai_spans=False, ) prior_response = "Of course! How can I assist you?" @@ -2401,13 +1740,13 @@ def original_invoke(self, *args, **kwargs): ] } - captured = capture_items("span") if span_streaming else capture_events() + captured = capture_items("span") with start_transaction(): wrapped_invoke = _wrap_pregel_invoke(original_invoke) wrapped_invoke(pregel, test_state) - data = _invoke_span_data(captured, span_streaming) + data = _invoke_span_data(captured) assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response @@ -2418,21 +1757,17 @@ def original_invoke(self, *args, **kwargs): assert data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-2025-04-14" -@pytest.mark.parametrize("span_streaming", [True, False]) def test_pregel_ainvoke_message_delta_ignores_gen_ai_inputs_setting( sentry_init, - capture_events, capture_items, - span_streaming, ): sentry_init( integrations=[LanggraphIntegration()], traces_sample_rate=1.0, - trace_lifecycle="stream" if span_streaming else "static", + trace_lifecycle="stream", _experiments={ "data_collection": {"gen_ai": {"inputs": False, "outputs": True}} }, - stream_gen_ai_spans=False, ) prior_response = "It is sunny in Berlin." @@ -2480,11 +1815,11 @@ async def run_test(): wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) return await wrapped_ainvoke(pregel, test_state) - captured = capture_items("span") if span_streaming else capture_events() + captured = capture_items("span") asyncio.run(run_test()) - data = _invoke_span_data(captured, span_streaming) + data = _invoke_span_data(captured) assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response @@ -2493,77 +1828,3 @@ async def run_test(): assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 assert data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 assert data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-2025-04-14" - - -@pytest.mark.parametrize( - "data_collection, send_default_pii, expect_available_tools", - [ - pytest.param( - {"gen_ai": {"inputs": True}}, - False, - True, - id="gen-ai-inputs-enabled-overrides-pii-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - True, - False, - id="gen-ai-inputs-disabled-overrides-pii-enabled", - ), - pytest.param( - {}, - False, - True, - id="gen-ai-omitted-defaults-to-enabled", - ), - pytest.param( - None, - False, - True, - id="no-data-collection-collects-regardless-of-pii", - ), - ], -) -def test_state_graph_compile_gates_available_tools_only_when_data_collection_configured( - sentry_init, - capture_events, - data_collection, - send_default_pii, - expect_available_tools, -): - init_kwargs = { - "integrations": [LanggraphIntegration()], - "traces_sample_rate": 1.0, - "send_default_pii": send_default_pii, - "stream_gen_ai_spans": False, - } - if data_collection is not None: - init_kwargs["_experiments"] = {"data_collection": data_collection} - - sentry_init(**init_kwargs) - - graph = MockStateGraph() - - def original_compile(self, *args, **kwargs): - return MockCompiledGraph(self.name) - - events = capture_events() - - with patch("sentry_sdk.integrations.langgraph.StateGraph"), start_transaction(): - wrapped_compile = _wrap_state_graph_compile(original_compile) - wrapped_compile(graph, model="test-model", checkpointer=None) - - tx = events[0] - agent_spans = [span for span in tx["spans"] if span["op"] == OP.GEN_AI_CREATE_AGENT] - assert len(agent_spans) == 1 - data = agent_spans[0]["data"] - - if expect_available_tools: - assert data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS] == [ - "search_tool", - "calculator", - ] - else: - assert SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS not in data - - assert data[SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" From fe50d42980fe4427e659d243fe4607ee581173f8 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 4 Sep 2026 15:46:48 +0200 Subject: [PATCH 2/2] remove start_transaction --- .../integrations/langgraph/test_langgraph.py | 113 +++++++----------- 1 file changed, 45 insertions(+), 68 deletions(-) diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index 76055f9379..224edcf7b5 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -10,7 +10,6 @@ from langgraph.errors import GraphBubbleUp import sentry_sdk -from sentry_sdk import start_transaction from sentry_sdk.consts import OP, SPANDATA @@ -205,9 +204,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) assert result is not None @@ -306,10 +304,9 @@ async def original_ainvoke(self, *args, **kwargs): return {"messages": new_messages} async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result items = capture_items("span") @@ -376,7 +373,7 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(), pytest.raises(Exception, match="Graph execution failed"): + with pytest.raises(Exception, match="Graph execution failed"): wrapped_invoke = _wrap_pregel_invoke(original_invoke) wrapped_invoke(pregel, test_state) @@ -412,9 +409,7 @@ async def original_ainvoke(self, *args, **kwargs): raise Exception("Async graph execution failed") async def run_error_test(): - with start_transaction(), pytest.raises( - Exception, match="Async graph execution failed" - ): + with pytest.raises(Exception, match="Async graph execution failed"): wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) await wrapped_ainvoke(pregel, test_state) @@ -459,9 +454,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, {"messages": []}) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, {"messages": []}) sentry_sdk.flush() spans = [item.payload for item in items] @@ -537,9 +531,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) assert result is not None @@ -618,10 +611,9 @@ async def original_ainvoke(self, *args, **kwargs): return {"messages": new_messages} async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result items = capture_items("span") @@ -715,9 +707,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) assert result is not None @@ -801,10 +792,9 @@ async def original_ainvoke(self, *args, **kwargs): return {"messages": new_messages} async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result items = capture_items("span") @@ -881,9 +871,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) assert result is not None @@ -958,10 +947,9 @@ async def original_ainvoke(self, *args, **kwargs): return {"messages": new_messages} async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result items = capture_items("span") @@ -1053,9 +1041,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) assert result is not None @@ -1142,10 +1129,9 @@ async def original_ainvoke(self, *args, **kwargs): return {"messages": new_messages} async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + result = await wrapped_ainvoke(pregel, test_state) + return result items = capture_items("span") @@ -1264,9 +1250,8 @@ def original_invoke(self, *args, **kwargs): items = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + result = wrapped_invoke(pregel, test_state) assert result is not None @@ -1330,14 +1315,11 @@ def __init__(self, content, message_type="human"): pregel = MockPregelInstance(compiled_graph) items = capture_items("span") - with start_transaction(name="langgraph tx"): - # Use the wrapped invoke function directly - from sentry_sdk.integrations.langgraph import _wrap_pregel_invoke + # Use the wrapped invoke function directly + from sentry_sdk.integrations.langgraph import _wrap_pregel_invoke - wrapped_invoke = _wrap_pregel_invoke( - lambda self, state_data: {"result": "success"} - ) - wrapped_invoke(pregel, state_data) + wrapped_invoke = _wrap_pregel_invoke(lambda self, state_data: {"result": "success"}) + wrapped_invoke(pregel, state_data) sentry_sdk.flush() span = next(item.payload for item in items) @@ -1474,9 +1456,8 @@ def original_invoke(self, *args, **kwargs): captured = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) data = _invoke_span_data(captured) @@ -1570,9 +1551,8 @@ def original_invoke(self, *args, **kwargs): captured = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) data = _invoke_span_data(captured) @@ -1662,9 +1642,8 @@ async def original_ainvoke(self, *args, **kwargs): } async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - return await wrapped_ainvoke(pregel, test_state) + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + return await wrapped_ainvoke(pregel, test_state) captured = capture_items("span") @@ -1742,9 +1721,8 @@ def original_invoke(self, *args, **kwargs): captured = capture_items("span") - with start_transaction(): - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) + wrapped_invoke = _wrap_pregel_invoke(original_invoke) + wrapped_invoke(pregel, test_state) data = _invoke_span_data(captured) @@ -1811,9 +1789,8 @@ async def original_ainvoke(self, *args, **kwargs): } async def run_test(): - with start_transaction(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - return await wrapped_ainvoke(pregel, test_state) + wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) + return await wrapped_ainvoke(pregel, test_state) captured = capture_items("span")