feat(tracing): emit token usage from LangGraph tracing handler

Author: levilentz

src/agentex/lib/adk/_modules/_langgraph_tracing.py

@@ -14,6 +14,7 @@
 from agentex.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.adk._modules.tracing import TracingModule
+from agentex.lib.core.tracing.usage import usage_from_counts
 
 logger = make_logger(__name__)
 
@@ -220,9 +221,44 @@ def _serialize_llm_result(response: LLMResult) -> dict[str, Any]:
 
             if msg and hasattr(msg, "tool_calls") and msg.tool_calls:
                 output["tool_calls"] = [{"name": tc["name"], "args": tc["args"]} for tc in msg.tool_calls]
+
+            # Per-call usage for billing; deduped against any turn aggregate
+            usage = _extract_usage(response, msg)
+            if usage:
+                output["usage"] = usage
     return output
 
 
+def _extract_usage(response: LLMResult, msg: Any) -> dict[str, Any] | None:
+    """Pull token usage from an LLMResult in backend-billable key spellings.
+
+    Prefers the provider-normalized ``AIMessage.usage_metadata``; falls back to
+    ``LLMResult.llm_output["token_usage"]`` (OpenAI-style keys, which the
+    backend also accepts).
+    """
+    usage_metadata = getattr(msg, "usage_metadata", None) if msg is not None else None
+    if isinstance(usage_metadata, dict) and usage_metadata:
+        input_details = usage_metadata.get("input_token_details") or {}
+        output_details = usage_metadata.get("output_token_details") or {}
+        return usage_from_counts(
+            input_tokens=usage_metadata.get("input_tokens"),
+            output_tokens=usage_metadata.get("output_tokens"),
+            total_tokens=usage_metadata.get("total_tokens"),
+            cached_input_tokens=input_details.get("cache_read"),
+            reasoning_tokens=output_details.get("reasoning"),
+        )
+
+    llm_output = getattr(response, "llm_output", None) or {}
+    token_usage = llm_output.get("token_usage") or llm_output.get("usage")
+    if isinstance(token_usage, dict) and token_usage:
+        return {
+            key: value
+            for key, value in token_usage.items()
+            if isinstance(value, (int, float))
+        }
+    return None
+
+
 def create_langgraph_tracing_handler(
     trace_id: str,
     parent_span_id: str | None = None,

tests/lib/adk/test_langgraph_tracing_usage.py

@@ -0,0 +1,91 @@
+"""Tests for usage extraction in the LangGraph tracing handler.
+
+Uses duck-typed fakes rather than real langchain_core objects because the ADK
+test conftest mocks langchain_core; ``_serialize_llm_result`` only relies on
+attribute access, so the fakes match the runtime contract.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+from agentex.lib.adk._modules._langgraph_tracing import _serialize_llm_result
+
+
+def _make_result(msg=None, text: str = "hello", llm_output=None):
+    generation = SimpleNamespace(text=text, message=msg)
+    return SimpleNamespace(generations=[[generation]], llm_output=llm_output)
+
+
+def _make_message(content="hello", usage_metadata=None, tool_calls=None):
+    return SimpleNamespace(content=content, usage_metadata=usage_metadata, tool_calls=tool_calls)
+
+
+class TestSerializeLLMResultUsage:
+    def test_usage_from_usage_metadata(self):
+        msg = _make_message(
+            usage_metadata={
+                "input_tokens": 100,
+                "output_tokens": 40,
+                "total_tokens": 140,
+                "input_token_details": {"cache_read": 25},
+                "output_token_details": {"reasoning": 12},
+            }
+        )
+        output = _serialize_llm_result(_make_result(msg=msg))
+
+        assert output["content"] == "hello"
+        assert output["usage"] == {
+            "input_tokens": 100,
+            "output_tokens": 40,
+            "total_tokens": 140,
+            "cached_input_tokens": 25,
+            "reasoning_tokens": 12,
+        }
+
+    def test_usage_metadata_without_details(self):
+        msg = _make_message(usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15})
+        output = _serialize_llm_result(_make_result(msg=msg))
+
+        assert output["usage"] == {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
+
+    def test_fallback_to_llm_output_token_usage(self):
+        msg = _make_message(usage_metadata=None)
+        result = _make_result(
+            msg=msg,
+            llm_output={
+                "token_usage": {
+                    "prompt_tokens": 30,
+                    "completion_tokens": 20,
+                    "total_tokens": 50,
+                    # Nested detail dicts are dropped: the span blob stays flat
+                    "completion_tokens_details": {"reasoning_tokens": 4},
+                }
+            },
+        )
+        output = _serialize_llm_result(result)
+
+        assert output["usage"] == {
+            "prompt_tokens": 30,
+            "completion_tokens": 20,
+            "total_tokens": 50,
+        }
+
+    def test_no_usage_available_omits_key(self):
+        msg = _make_message(usage_metadata=None)
+        output = _serialize_llm_result(_make_result(msg=msg))
+
+        assert "usage" not in output
+        assert output["content"] == "hello"
+
+    def test_content_blocks_and_tool_calls_still_serialized(self):
+        msg = _make_message(
+            content=[{"type": "text", "text": "block text"}],
+            usage_metadata={"input_tokens": 1, "output_tokens": 2, "total_tokens": 3},
+            tool_calls=[{"name": "search", "args": {"q": "x"}}],
+        )
+        output = _serialize_llm_result(_make_result(msg=msg))
+
+        assert output["content"] == "block text"
+        assert output["tool_calls"] == [{"name": "search", "args": {"q": "x"}}]
+        assert output["usage"]["total_tokens"] == 3