From 71cefa2f8773ada48cdb4c1ce30d19059a2e5100 Mon Sep 17 00:00:00 2001 From: "sentry-junior[bot]" <264270552+sentry-junior[bot]@users.noreply.github.com> Date: Fri, 13 Nov 2026 11:48:13 +0000 Subject: [PATCH] fix(langgraph): Create invoke_agent span for stream() and astream() Pregel.invoke and Pregel.ainvoke run through Pregel.stream and Pregel.astream, not the other way round, so graphs run with stream() or astream() got no invoke_agent span. Wrap the stream methods too, and skip the inner stream call of an invoke so it still creates exactly one span. Fixes #7871 Co-Authored-By: Johannes Daxboeck <114897+dingsdax@users.noreply.github.com> --- sentry_sdk/integrations/langgraph.py | 323 ++++++++---------- .../integrations/langgraph/test_langgraph.py | 132 +++++++ 2 files changed, 280 insertions(+), 175 deletions(-) diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index af287b0fd4..31b27d0b6f 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -1,5 +1,6 @@ +from contextlib import contextmanager from functools import wraps -from typing import Any, Callable, List, Optional +from typing import Any, AsyncIterator, Callable, Iterator, List, Optional, Tuple import sentry_sdk from sentry_sdk.ai.utils import ( @@ -19,7 +20,7 @@ has_span_streaming_enabled, should_truncate_gen_ai_input, ) -from sentry_sdk.utils import has_data_collection_enabled, safe_serialize +from sentry_sdk.utils import ContextVar, has_data_collection_enabled, safe_serialize try: from langgraph.errors import GraphBubbleUp @@ -28,6 +29,11 @@ except ImportError: raise DidNotEnable("langgraph not installed") +# The graph whose invoke() or ainvoke() is running. LangGraph runs invoke() through +# stream() and ainvoke() through astream(), so the stream wrappers skip that inner +# call instead of creating a second span. +_invoking_graph = ContextVar("sentry_langgraph_invoking_graph", default=None) + class LanggraphIntegration(Integration): identifier = "langgraph" @@ -42,15 +48,17 @@ def setup_once() -> None: # LangGraph lets users create agents using a StateGraph or the Functional API. # StateGraphs are then compiled to a CompiledStateGraph. Both CompiledStateGraph and # the functional API execute on a Pregel instance. Pregel is the runtime for the graph - # and the invocation happens on Pregel, so patching the invoke methods takes care of both. - # 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. + # and the invocation happens on Pregel, so patching its invoke and stream methods + # takes care of both. StateGraph.compile = _wrap_state_graph_compile(StateGraph.compile) if hasattr(Pregel, "invoke"): Pregel.invoke = _wrap_pregel_invoke(Pregel.invoke) if hasattr(Pregel, "ainvoke"): Pregel.ainvoke = _wrap_pregel_ainvoke(Pregel.ainvoke) + if hasattr(Pregel, "stream"): + Pregel.stream = _wrap_pregel_stream(Pregel.stream) + if hasattr(Pregel, "astream"): + Pregel.astream = _wrap_pregel_astream(Pregel.astream) def _should_record_inputs(integration: "LanggraphIntegration") -> bool: @@ -184,104 +192,82 @@ def new_compile(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": return new_compile -def _wrap_pregel_invoke(f: "Callable[..., Any]") -> "Callable[..., Any]": - @wraps(f) - def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - client = sentry_sdk.get_client() - integration = client.get_integration(LanggraphIntegration) - if integration is None: - return f(self, *args, **kwargs) +@contextmanager +def _invoke_agent_span( + graph: "Any", args: "Any", integration: "LanggraphIntegration" +) -> "Iterator[Tuple[Any, Optional[List[Any]]]]": + """Start an invoke_agent span for one graph run and set its request data.""" + client = sentry_sdk.get_client() + graph_name = _get_graph_name(graph) + span_name = f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent" - 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): + span_context: "Any" = sentry_sdk.traces.start_span( + name=span_name, + attributes={ + "sentry.op": OP.GEN_AI_INVOKE_AGENT, + "sentry.origin": LanggraphIntegration.origin, + }, + ) + else: + span_context = get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name=span_name, + origin=LanggraphIntegration.origin, ) - 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 should_truncate_gen_ai_input(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, - ) + with span_context as span: + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) + set_on_span(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + if graph_name: + set_on_span(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name) + set_on_span(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 should_truncate_gen_ai_input(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) + yield span, input_messages - _set_response_attributes(span, input_messages, result, integration) - 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 - ) - if should_truncate_gen_ai_input(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, - ) +def _wrap_pregel_invoke(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + if integration is None: + return f(self, *args, **kwargs) + with _invoke_agent_span(self, args, integration) as (span, input_messages): + token = _invoking_graph.set(self) + try: result = f(self, *args, **kwargs) + finally: + _invoking_graph.reset(token) - _set_response_attributes(span, input_messages, result, integration) + _set_response_attributes(span, input_messages, result, integration) - return result + return result return new_invoke @@ -289,99 +275,86 @@ def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": def _wrap_pregel_ainvoke(f: "Callable[..., Any]") -> "Callable[..., Any]": @wraps(f) async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": - client = sentry_sdk.get_client() - integration = client.get_integration(LanggraphIntegration) + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) if integration is None: 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" - ) + with _invoke_agent_span(self, args, integration) as (span, input_messages): + token = _invoking_graph.set(self) + try: + result = await f(self, *args, **kwargs) + finally: + _invoking_graph.reset(token) - 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 should_truncate_gen_ai_input(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, - ) + _set_response_attributes(span, input_messages, result, integration) - result = await f(self, *args, **kwargs) + return result - _set_response_attributes(span, input_messages, result, integration) + return new_ainvoke - return result - 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") - - 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 should_truncate_gen_ai_input(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, - ) +def _wrap_pregel_stream(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_stream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + # Pregel.invoke runs through Pregel.stream, and the invoke wrapper already + # created the span for this run. + if integration is None or _invoking_graph.get() is self: + return f(self, *args, **kwargs) - result = await f(self, *args, **kwargs) + def stream_with_span() -> "Iterator[Any]": + with _invoke_agent_span(self, args, integration) as ( + span, + input_messages, + ): + stream = f(self, *args, **kwargs) + last_chunk = None + try: + for chunk in stream: + last_chunk = chunk + yield chunk + except GeneratorExit: + # The caller stopped iterating early, which isn't an error. + stream.close() + return - _set_response_attributes(span, input_messages, result, integration) + _set_response_attributes(span, input_messages, last_chunk, integration) - return result + return stream_with_span() - return new_ainvoke + return new_stream + + +def _wrap_pregel_astream(f: "Callable[..., Any]") -> "Callable[..., Any]": + @wraps(f) + def new_astream(self: "Any", *args: "Any", **kwargs: "Any") -> "Any": + integration = sentry_sdk.get_client().get_integration(LanggraphIntegration) + # Pregel.ainvoke runs through Pregel.astream, and the ainvoke wrapper already + # created the span for this run. + if integration is None or _invoking_graph.get() is self: + return f(self, *args, **kwargs) + + async def astream_with_span() -> "AsyncIterator[Any]": + with _invoke_agent_span(self, args, integration) as ( + span, + input_messages, + ): + stream = f(self, *args, **kwargs) + last_chunk = None + try: + async for chunk in stream: + last_chunk = chunk + yield chunk + except GeneratorExit: + # The caller stopped iterating early, which isn't an error. + await stream.aclose() + return + + _set_response_attributes(span, input_messages, last_chunk, integration) + + return astream_with_span() + + return new_astream def _get_new_messages( diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index 7b6a6cc3d9..3d39f7cc35 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -38,7 +38,9 @@ def mock_langgraph_imports(): LanggraphIntegration, _parse_langgraph_messages, _wrap_pregel_ainvoke, + _wrap_pregel_astream, _wrap_pregel_invoke, + _wrap_pregel_stream, _wrap_state_graph_compile, ) @@ -2655,3 +2657,133 @@ def original_compile(self, *args, **kwargs): assert SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS not in data assert data[SPANDATA.GEN_AI_AGENT_NAME] == "test_graph" + + +class MockStreamingPregel: + """Like Pregel: invoke() and ainvoke() consume stream() and astream().""" + + def __init__(self, name="stream_graph"): + self.name = name + self.stream_closed = False + + def _chunks(self, state): + # stream_mode="values" yields the full state after each step + return [ + state, + { + "messages": state["messages"] + + [MockMessage("Streamed response", name="assistant")] + }, + ] + + def stream(self, state, config=None): + try: + yield from self._chunks(state) + finally: + self.stream_closed = True + + async def astream(self, state, config=None): + try: + for chunk in self._chunks(state): + yield chunk + finally: + self.stream_closed = True + + def invoke(self, state, config=None): + latest = None + for chunk in self.stream(state, config): + latest = chunk + return latest + + async def ainvoke(self, state, config=None): + latest = None + async for chunk in self.astream(state, config): + latest = chunk + return latest + + +MockStreamingPregel.stream = _wrap_pregel_stream(MockStreamingPregel.stream) +MockStreamingPregel.astream = _wrap_pregel_astream(MockStreamingPregel.astream) +MockStreamingPregel.invoke = _wrap_pregel_invoke(MockStreamingPregel.invoke) +MockStreamingPregel.ainvoke = _wrap_pregel_ainvoke(MockStreamingPregel.ainvoke) + + +def _run_graph(graph, method, state, stop_after=None): + """Run graph.() in a transaction, consuming at most stop_after chunks.""" + + async def run_async(): + with start_transaction(): + if method == "ainvoke": + await graph.ainvoke(state) + return + stream = graph.astream(state) + seen = 0 + async for _ in stream: + seen += 1 + if seen == stop_after: + break + await stream.aclose() + + if method in ("ainvoke", "astream"): + asyncio.run(run_async()) + return + + with start_transaction(): + if method == "invoke": + graph.invoke(state) + return + stream = graph.stream(state) + seen = 0 + for _ in stream: + seen += 1 + if seen == stop_after: + break + stream.close() + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("method", ["stream", "astream", "invoke", "ainvoke"]) +def test_pregel_stream_and_invoke_create_one_invoke_agent_span( + sentry_init, capture_events, capture_items, method, span_streaming +): + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + stream_gen_ai_spans=span_streaming, + trace_lifecycle="stream" if span_streaming else "static", + ) + captured = capture_items("span") if span_streaming else capture_events() + + graph = MockStreamingPregel("stream_graph") + _run_graph(graph, method, {"messages": [MockMessage("Hi", name="user")]}) + + # _invoke_span_data asserts there's exactly one invoke_agent span + data = _invoke_span_data(captured, span_streaming) + assert data[SPANDATA.GEN_AI_AGENT_NAME] == "stream_graph" + assert data[SPANDATA.GEN_AI_OPERATION_NAME] == "invoke_agent" + assert SPANDATA.GEN_AI_REQUEST_MESSAGES in data + assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == "Streamed response" + + +@pytest.mark.parametrize("method", ["stream", "astream"]) +def test_pregel_stream_closed_early(sentry_init, capture_events, method): + sentry_init( + integrations=[LanggraphIntegration()], + traces_sample_rate=1.0, + send_default_pii=True, + stream_gen_ai_spans=False, + ) + events = capture_events() + + graph = MockStreamingPregel("stream_graph") + _run_graph( + graph, method, {"messages": [MockMessage("Hi", name="user")]}, stop_after=1 + ) + + assert graph.stream_closed + (tx,) = events + (span,) = [s for s in tx["spans"] if s["op"] == OP.GEN_AI_INVOKE_AGENT] + # Stopping early isn't an error + assert span.get("status") != "internal_error" + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"]