[BUG]: ddtrace >= 3.5.0 breaks langgraph integration when using astream_events

[vhiairrassary]

Tracer Version(s)

3.5.0 and later

Python Version(s)

Python 3.13.3 (but I can reproduce it with 3.11.12 as well)

Pip Version(s)

uv 0.7.2 (Homebrew 2025-04-30)

Bug Report

ddtrace >= 3.5.0 breaks langgraph integration when using astream_events (3.4.2 is OK), with the following error:

Traceback (most recent call last):
  File "/Users/vhiairrassary/abc/test.py", line 30, in <module>
    event_loop.run_until_complete(main())
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^
  File "/Users/vhiairrassary/.local/share/uv/python/cpython-3.13.3-macos-aarch64-none/lib/python3.13/asyncio/base_events.py", line 719, in run_until_complete
    return future.result()
           ~~~~~~~~~~~~~^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/contrib/internal/asyncio/patch.py", line 50, in traced_coro
    return await coro
           ^^^^^^^^^^
  File "/Users/vhiairrassary/abc/test.py", line 23, in main
    async for event in agent.astream_events(agent_input, version="v2", config=config):
        print(f"event: {event}")
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/langchain_core/runnables/base.py", line 1394, in astream_events
    async for event in event_stream:
        yield event
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/langchain_core/tracers/event_stream.py", line 1021, in _astream_events_implementation_v2
    await task
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/contrib/internal/asyncio/patch.py", line 50, in traced_coro
    return await coro
           ^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/langchain_core/tracers/event_stream.py", line 976, in consume_astream
    async for _ in event_streamer.tap_output_aiter(run_id, stream):
        # All the content will be picked up
        pass
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/langchain_core/tracers/event_stream.py", line 181, in tap_output_aiter
    first = await py_anext(output, default=sentinel)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/langchain_core/utils/aiter.py", line 78, in anext_impl
    return await __anext__(iterator)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/contrib/internal/langgraph/patch.py", line 174, in _astream
    item = await result.__anext__()
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/langgraph/pregel/__init__.py", line 2758, in astream
    while loop.tick(input_keys=self.input_channels):
          ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/contrib/internal/trace_utils.py", line 329, in wrapper
    return func(mod, pin, wrapped, instance, args, kwargs)
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/contrib/internal/langgraph/patch.py", line 205, in patched_pregel_loop_tick
    integration.llmobs_handle_pregel_loop_tick(finished_tasks, next_tasks, result, is_subgraph_node)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/llmobs/_integrations/langgraph.py", line 76, in llmobs_handle_pregel_loop_tick
    self._handle_finished_graph(graph_span, finished_tasks, is_subgraph_node)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vhiairrassary/abc/.venv/lib/python3.13/site-packages/ddtrace/llmobs/_integrations/langgraph.py", line 91, in _handle_finished_graph
    {**self._graph_nodes_by_task_id[task_id]["span"], "attributes": {"from": "output", "to": "output"}}
       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^
KeyError: 'span'

Reproduction Code

import asyncio
import uuid
from langchain.chat_models import init_chat_model
from langgraph.prebuilt import create_react_agent
from langchain_core.runnables import RunnableConfig


async def main():
    model = init_chat_model(
        "anthropic.claude-3-sonnet-20240229-v1:0",
        model_provider="bedrock_converse",
    )
    agent = create_react_agent(model, [])

    config = RunnableConfig(
        configurable={
            "thread_id": uuid.uuid4(),
        }
    )

    agent_input = {"messages": "How much is 10 + 4"}

    async for event in agent.astream_events(agent_input, version="v2", config=config):
        print(f"event: {event}")


if __name__ == "__main__":
    event_loop = asyncio.new_event_loop()
    asyncio.set_event_loop(event_loop)
    event_loop.run_until_complete(main())

Error Logs

No response

Libraries in Use

• pyproject.toml

[project]
name = "debug-ddtrace"
version = "0.0.1"
requires-python = "===3.13.3"
dependencies = [
    "ddtrace===3.5.1",
    #"ddtrace===3.4.2",
    "langchain>=0.3.25",
    "langchain-aws>=0.2.23",
    "langgraph>=0.4.3"
]

• .python-version

3.13.3

Operating System

Darwin Victor-Hiairrassary 24.4.0 Darwin Kernel Version 24.4.0: Fri Apr 11 18:33:47 PDT 2025; root:xnu-11417.101.15~117/RELEASE_ARM64_T6000 arm64

[sabrenner]

Hi @vhiairrassary, thanks for opening this issue, and for providing a reproduction! I am able to reproduce the issue - it seems using astream instead of astream_events works as a potential workaround for now but I am investigating what is causing the issue and will link a fix PR to this issue once I have one up!