fix(deepagents): capture root agent input messages

Author: sipercai

instrumentation-loongsuite/loongsuite-instrumentation-deepagents/src/opentelemetry/instrumentation/deepagents/internal/patch.py

@@ -30,6 +30,7 @@
 from collections.abc import AsyncIterator, Iterator
 from contextlib import suppress
 from typing import Any, Callable
+from weakref import WeakSet
 
 from wrapt import ObjectProxy, wrap_function_wrapper
 
@@ -48,6 +49,8 @@
 _MISSING = object()
 _top_level_original: Any = _MISSING
 _top_level_patched = False
+_wrapped_graphs = WeakSet()
+_strong_wrapped_graphs: list[Any] = []
 
 
 def _create_deep_agent_wrapper(
@@ -109,6 +112,7 @@ def uninstrument_create_deep_agent() -> None:
     with suppress(Exception):
         module = importlib.import_module(CREATE_DEEP_AGENT_MODULE)
         unwrap(module, CREATE_DEEP_AGENT_NAME)
+    _restore_wrapped_graph_methods()
     _restore_top_level_create_deep_agent()
 
 
@@ -186,7 +190,7 @@ def _wrap_graph_methods(graph: Any) -> None:
         return
 
     originals: dict[str, Callable[..., Any]] = {}
-    for method_name in ("invoke", "ainvoke", "stream", "astream"):
+    for method_name in ("invoke", "ainvoke", "stream", "astream", "with_config"):
         original = getattr(graph, method_name, None)
         if original is None:
             continue
@@ -204,6 +208,38 @@ def _wrap_graph_methods(graph: Any) -> None:
     with suppress(Exception):
         setattr(graph, GRAPH_ORIGINAL_METHODS_ATTR, originals)
         setattr(graph, GRAPH_METHODS_WRAPPED_ATTR, True)
+    _track_wrapped_graph(graph)
+
+
+def _track_wrapped_graph(graph: Any) -> None:
+    try:
+        _wrapped_graphs.add(graph)
+    except TypeError:
+        _strong_wrapped_graphs.append(graph)
+
+
+def _restore_wrapped_graph_methods() -> None:
+    for graph in [*list(_wrapped_graphs), *_strong_wrapped_graphs]:
+        _restore_graph_methods(graph)
+    _wrapped_graphs.clear()
+    _strong_wrapped_graphs.clear()
+
+
+def _restore_graph_methods(graph: Any) -> None:
+    originals = getattr(graph, GRAPH_ORIGINAL_METHODS_ATTR, None)
+    if isinstance(originals, dict):
+        for method_name, original in originals.items():
+            with suppress(Exception):
+                setattr(graph, method_name, original)
+
+    for attr_name in (
+        GRAPH_ORIGINAL_METHODS_ATTR,
+        GRAPH_METHODS_WRAPPED_ATTR,
+        REACT_AGENT_METADATA_KEY,
+        DEEPAGENTS_METADATA_KEY,
+    ):
+        with suppress(Exception):
+            delattr(graph, attr_name)
 
 
 def _make_method_wrapper(
@@ -222,21 +258,28 @@ async def ainvoke_wrapper(*args: Any, **kwargs: Any) -> Any:
 
         def stream_wrapper(*args: Any, **kwargs: Any) -> Iterator[Any]:
             args, kwargs = _rewrite_config(args, kwargs)
-            yield from original(*args, **kwargs)
+            return original(*args, **kwargs)
 
         return stream_wrapper
 
     if method_name == "astream":
 
-        async def astream_wrapper(
-            *args: Any, **kwargs: Any
-        ) -> AsyncIterator[Any]:
+        def astream_wrapper(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
             args, kwargs = _rewrite_config(args, kwargs)
-            async for chunk in original(*args, **kwargs):
-                yield chunk
+            return original(*args, **kwargs)
 
         return astream_wrapper
 
+    if method_name == "with_config":
+
+        def with_config_wrapper(*args: Any, **kwargs: Any) -> Any:
+            graph = original(*args, **kwargs)
+            _mark_graph(graph)
+            _wrap_graph_methods(graph)
+            return graph
+
+        return with_config_wrapper
+
     def invoke_wrapper(*args: Any, **kwargs: Any) -> Any:
         args, kwargs = _rewrite_config(args, kwargs)
         return original(*args, **kwargs)

instrumentation-loongsuite/loongsuite-instrumentation-deepagents/tests/test_deepagents_spans.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import json
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Optional
 
@@ -172,6 +173,15 @@ def test_deepagents_root_span_is_agent_and_single_step(
 
     spans = span_exporter.get_finished_spans()
     agent_span = _assert_single_agent(spans, "deep_test_agent")
+    input_raw = agent_span.attributes.get("gen_ai.input.messages")
+    assert input_raw, "AGENT span missing gen_ai.input.messages"
+    input_messages = json.loads(input_raw)
+    assert input_messages == [
+        {
+            "role": "user",
+            "parts": [{"content": "hello", "type": "text"}],
+        }
+    ]
 
     step_spans = _spans_by_kind(spans, "STEP")
     assert len(step_spans) == 1

instrumentation-loongsuite/loongsuite-instrumentation-deepagents/tests/test_instrumentor.py

@@ -17,10 +17,50 @@
 from __future__ import annotations
 
 from opentelemetry.instrumentation.deepagents import DeepAgentsInstrumentor
+from opentelemetry.instrumentation.deepagents.internal.patch import (
+    DEEPAGENTS_METADATA_KEY,
+    GRAPH_METHODS_WRAPPED_ATTR,
+    GRAPH_ORIGINAL_METHODS_ATTR,
+    REACT_AGENT_METADATA_KEY,
+    _mark_graph,
+    _wrap_graph_methods,
+    uninstrument_create_deep_agent,
+)
 from opentelemetry.instrumentation.langchain import LangChainInstrumentor
 from opentelemetry.instrumentation.langgraph import LangGraphInstrumentor
 
 
+class _DummyGraph:
+    def __init__(self):
+        self.seen_stream_config = None
+        self.bound_config = None
+
+    def _effective_config(self, config):
+        if self.bound_config is None:
+            return config
+        if config is None:
+            return self.bound_config
+
+        merged = {**self.bound_config, **config}
+        metadata = dict(self.bound_config.get("metadata") or {})
+        metadata.update(config.get("metadata") or {})
+        merged["metadata"] = metadata
+        return merged
+
+    def invoke(self, _input, config=None):
+        return self._effective_config(config)
+
+    def stream(self, _input, config=None):
+        effective_config = self._effective_config(config)
+        self.seen_stream_config = effective_config
+        return iter([effective_config])
+
+    def with_config(self, config=None, **_kwargs):
+        graph = _DummyGraph()
+        graph.bound_config = config
+        return graph
+
+
 def _cleanup_instrumentors() -> None:
     for instrumentor in (
         DeepAgentsInstrumentor(),
@@ -50,6 +90,57 @@ def test_uninstrument_cleans_dependencies_it_instrumented():
         _cleanup_instrumentors()
 
 
+def test_uninstrument_restores_wrapped_graph_methods():
+    _cleanup_instrumentors()
+    graph = _DummyGraph()
+
+    try:
+        _mark_graph(graph)
+        _wrap_graph_methods(graph)
+
+        assert getattr(graph, GRAPH_METHODS_WRAPPED_ATTR) is True
+        assert getattr(graph, REACT_AGENT_METADATA_KEY) is True
+        assert getattr(graph, DEEPAGENTS_METADATA_KEY) is True
+
+        invoke_config = graph.invoke("hello")
+        assert invoke_config["metadata"][DEEPAGENTS_METADATA_KEY] is True
+
+        stream_iter = graph.stream("hello")
+        assert (
+            graph.seen_stream_config["metadata"][DEEPAGENTS_METADATA_KEY]
+            is True
+        )
+        assert next(stream_iter)["metadata"][DEEPAGENTS_METADATA_KEY] is True
+
+        child_graph = graph.with_config({"metadata": {"customer": "kept"}})
+        assert getattr(child_graph, GRAPH_METHODS_WRAPPED_ATTR) is True
+        assert getattr(child_graph, DEEPAGENTS_METADATA_KEY) is True
+        child_config = child_graph.invoke("hello")
+        assert child_config["metadata"]["customer"] == "kept"
+        assert child_config["metadata"][DEEPAGENTS_METADATA_KEY] is True
+
+        uninstrument_create_deep_agent()
+
+        assert not hasattr(graph, GRAPH_METHODS_WRAPPED_ATTR)
+        assert not hasattr(graph, GRAPH_ORIGINAL_METHODS_ATTR)
+        assert not hasattr(graph, REACT_AGENT_METADATA_KEY)
+        assert not hasattr(graph, DEEPAGENTS_METADATA_KEY)
+        assert not hasattr(child_graph, GRAPH_METHODS_WRAPPED_ATTR)
+        assert not hasattr(child_graph, GRAPH_ORIGINAL_METHODS_ATTR)
+        assert not hasattr(child_graph, REACT_AGENT_METADATA_KEY)
+        assert not hasattr(child_graph, DEEPAGENTS_METADATA_KEY)
+        assert graph.invoke("hello") is None
+        assert child_graph.invoke("hello") == {"metadata": {"customer": "kept"}}
+
+        graph.seen_stream_config = "not-called"
+        stream_iter = graph.stream("hello")
+        assert graph.seen_stream_config is None
+        assert next(stream_iter) is None
+    finally:
+        uninstrument_create_deep_agent()
+        _cleanup_instrumentors()
+
+
 def test_uninstrument_preserves_preinstrumented_dependencies():
     _cleanup_instrumentors()
     langchain_instrumentor = LangChainInstrumentor()

instrumentation-loongsuite/loongsuite-instrumentation-langchain/CHANGELOG.md

@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Capture OpenAI-style LangGraph and DeepAgents message dictionaries as
+  `gen_ai.input.messages` on root `AGENT` spans.
 - Preserve structured LangChain tool call arguments when tool inputs do not use
   `input` or `query`, which keeps DeepAgents filesystem `read_file` arguments
   such as `file_path` available for skill-load telemetry.

instrumentation-loongsuite/loongsuite-instrumentation-langchain/src/opentelemetry/instrumentation/langchain/internal/_tracer.py

@@ -96,6 +96,7 @@
     LLMInvocation,
     OutputMessage,
     Text,
+    ToolCall,
 )
 from opentelemetry.util.genai.utils import (
     ContentCapturingMode,
@@ -957,6 +958,67 @@ def _extract_langgraph_input_message(msg: Any) -> InputMessage | None:
             return InputMessage(role=str(role), parts=[Text(content=content)])
         return None
 
+    # OpenAI-style message dict: {"role": "user", "content": "hello"}.
+    if isinstance(msg, dict):
+        content = msg.get("content")
+        if content is None:
+            content = msg.get("text")
+
+        parts: list[Any] = []
+        if isinstance(content, str) and content:
+            parts.append(Text(content=content))
+        elif isinstance(content, list):
+            for part in content:
+                if isinstance(part, str) and part:
+                    parts.append(Text(content=part))
+                elif isinstance(part, dict):
+                    text = part.get("text")
+                    if text is None:
+                        text = part.get("content")
+                    if isinstance(text, str) and text:
+                        parts.append(Text(content=text))
+
+        for tool_call in msg.get("tool_calls") or []:
+            if isinstance(tool_call, dict):
+                function = tool_call.get("function") or {}
+                if not isinstance(function, dict):
+                    function = {}
+                arguments = tool_call.get("args")
+                if arguments is None:
+                    arguments = tool_call.get("arguments")
+                if arguments is None:
+                    arguments = function.get("arguments", {})
+                parts.append(
+                    ToolCall(
+                        name=tool_call.get("name")
+                        or function.get("name")
+                        or "",
+                        arguments=arguments,
+                        id=tool_call.get("id"),
+                    )
+                )
+
+        if parts:
+            role = msg.get("role") or msg.get("type") or "user"
+            role_name = str(role).lower()
+            if role_name.endswith("_message"):
+                role_name = role_name[: -len("_message")]
+            elif role_name.endswith("message"):
+                role_name = role_name[: -len("message")]
+            role_map = {
+                "human": "user",
+                "ai": "assistant",
+                "system": "system",
+                "function": "tool",
+                "tool": "tool",
+                "chat": "user",
+            }
+            return InputMessage(
+                role=role_map.get(role_name, role_name),
+                parts=parts,
+            )
+        return None
+
     # LangChain message object (HumanMessage, AIMessage, etc.)
     content = getattr(msg, "content", None)
     if content and isinstance(content, str):

instrumentation-loongsuite/loongsuite-instrumentation-langchain/tests/test_agent_spans.py

@@ -18,6 +18,7 @@
 
 from opentelemetry.instrumentation.langchain.internal._tracer import (
     LoongsuiteTracer,
+    _extract_langgraph_input_message,
 )
 from opentelemetry.instrumentation.langchain.internal._utils import (
     AGENT_RUN_NAMES,
@@ -77,6 +78,45 @@ def test_agent_run_names_immutable(self):
         assert isinstance(AGENT_RUN_NAMES, frozenset)
 
 
+def test_extract_langgraph_input_message_from_openai_style_dict():
+    message = _extract_langgraph_input_message(
+        {"role": "user", "content": "hello"}
+    )
+
+    assert message is not None
+    assert message.role == "user"
+    assert len(message.parts) == 1
+    assert message.parts[0].type == "text"
+    assert message.parts[0].content == "hello"
+
+
+def test_extract_langgraph_input_message_keeps_empty_content_tool_call():
+    message = _extract_langgraph_input_message(
+        {
+            "role": "assistant",
+            "content": "",
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "type": "function",
+                    "function": {
+                        "name": "search",
+                        "arguments": '{"query":"hello"}',
+                    },
+                }
+            ],
+        }
+    )
+
+    assert message is not None
+    assert message.role == "assistant"
+    assert len(message.parts) == 1
+    assert message.parts[0].type == "tool_call"
+    assert message.parts[0].name == "search"
+    assert message.parts[0].arguments == '{"query":"hello"}'
+    assert message.parts[0].id == "call_1"
+
+
 def test_agent_context_colors_child_llm_and_tool_spans(
     tracer_provider, span_exporter
 ):