fix(langgraph): accumulate multi-step usage in LangGraphTurn [greptile]

_capture overwrote self._usage on every AIMessage, so a multi-step turn (text
-> tool decision -> final text) reported only the last LLM call's tokens and
silently dropped the rest — undercounting in any billing/monitoring that reads
turn.usage(). Accumulate additively across calls via _accumulate_turn_usage
(None+None stays None; real 0 contributes 0). Add a test asserting summed
input/output/total/cache/reasoning tokens across two AIMessages.

The separate 06-18 "TurnResult.usage empty via auto_send_turn" comment is
resolved by the foundation (emitter reads turn.usage() after stream exhaustion).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: declan-scale

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

@@ -62,6 +62,33 @@ def langgraph_usage_to_turn_usage(usage_metadata: Any, model: str | None) -> Tur
     )
 
 
+def _add_optional(a: int | None, b: int | None) -> int | None:
+    """Sum two optional token counts; ``None`` means 'not reported' on that side.
+
+    ``None + None`` stays ``None`` (model never reported usage), while a real 0
+    contributes 0 (preserving zero counts rather than coercing them away).
+    """
+    if a is None and b is None:
+        return None
+    return (a or 0) + (b or 0)
+
+
+def _accumulate_turn_usage(acc: TurnUsage, call: TurnUsage, model: str | None) -> TurnUsage:
+    """Add a single LLM call's usage into the running per-turn total.
+
+    A LangGraph turn can make multiple LLM calls (e.g. text -> tool decision ->
+    final text); summing them avoids silently dropping all but the last call.
+    """
+    return TurnUsage(
+        model=model,
+        input_tokens=_add_optional(acc.input_tokens, call.input_tokens),
+        output_tokens=_add_optional(acc.output_tokens, call.output_tokens),
+        total_tokens=_add_optional(acc.total_tokens, call.total_tokens),
+        cached_input_tokens=_add_optional(acc.cached_input_tokens, call.cached_input_tokens),
+        reasoning_tokens=_add_optional(acc.reasoning_tokens, call.reasoning_tokens),
+    )
+
+
 class LangGraphTurn:
     """HarnessTurn wrapping a LangGraph ``astream()`` event stream.
 
@@ -89,7 +116,8 @@ class LangGraphTurn:
     option is needed.
 
     Usage data is captured lazily via the ``on_final_ai_message`` callback and
-    is only valid after ``events`` has been fully consumed.
+    is only valid after ``events`` has been fully consumed. Multi-step turns
+    (more than one LLM call) accumulate usage additively across calls.
     """
 
     def __init__(self, stream: Any, model: str | None = None) -> None:
@@ -105,15 +133,20 @@ async def _generate_events(self) -> AsyncGenerator[StreamTaskMessage, None]:
         def _capture(ai_msg: Any) -> None:
             usage_metadata = getattr(ai_msg, "usage_metadata", None)
             if usage_metadata is not None:
-                self._usage = langgraph_usage_to_turn_usage(usage_metadata, self._model)
+                call_usage = langgraph_usage_to_turn_usage(usage_metadata, self._model)
+                # Accumulate across LLM calls — the callback fires once per agent
+                # node invocation, so a multi-step turn reports usage more than
+                # once; overwriting would drop all but the last call.
+                self._usage = _accumulate_turn_usage(self._usage, call_usage, self._model)
 
         async for ev in convert_langgraph_to_agentex_events(self._stream, on_final_ai_message=_capture):
             yield ev
 
     def usage(self) -> TurnUsage:
-        """Return the usage captured from the last AIMessage in the stream.
+        """Return the usage accumulated across all AIMessages in the stream.
 
-        Valid only after ``events`` has been fully consumed.
-        Returns a zero-usage ``TurnUsage`` if the model did not report usage.
+        Multi-step turns sum each LLM call's usage. Valid only after ``events``
+        has been fully consumed. Returns a zero-usage ``TurnUsage`` if the model
+        did not report usage.
         """
         return self._usage

tests/lib/adk/test_langgraph_turn.py

@@ -168,6 +168,47 @@ async def test_usage_captured_from_ai_message(self):
         assert usage.total_tokens == 15
         assert usage.model == "gpt-4"
 
+    async def test_usage_accumulates_across_multiple_ai_messages(self):
+        """A multi-step turn (>1 LLM call) sums usage instead of keeping only the last."""
+        from langchain_core.messages import AIMessage
+
+        first = AIMessage(
+            content="thinking",
+            usage_metadata={
+                "input_tokens": 10,
+                "output_tokens": 5,
+                "total_tokens": 15,
+                "input_token_details": {"cache_read": 2},
+                "output_token_details": {"reasoning": 1},
+            },
+        )
+        second = AIMessage(
+            content="answer",
+            usage_metadata={
+                "input_tokens": 20,
+                "output_tokens": 7,
+                "total_tokens": 27,
+                "input_token_details": {"cache_read": 3},
+                "output_token_details": {"reasoning": 4},
+            },
+        )
+        stream = _make_stream(
+            [
+                ("updates", {"agent": {"messages": [first]}}),
+                ("updates", {"agent": {"messages": [second]}}),
+            ]
+        )
+        turn = LangGraphTurn(stream, model="gpt-4")
+        await _drain(turn)
+
+        usage = turn.usage()
+        assert usage.input_tokens == 30
+        assert usage.output_tokens == 12
+        assert usage.total_tokens == 42
+        assert usage.cached_input_tokens == 5
+        assert usage.reasoning_tokens == 5
+        assert usage.model == "gpt-4"
+
     async def test_usage_not_updated_when_no_usage_metadata(self):
         from langchain_core.messages import AIMessage