refactor(adk): record_usage accepts harness TurnUsage; drop bespoke usage helpers

TurnSpan.record_usage now takes the TurnUsage every harness turn adapter reports (cost_usd lifted to data automatically) or a plain dict, replacing the individual-count kwargs and the lib/core/tracing/usage.py helpers that duplicated what next's harness provides.

Author: levilentz

examples/tutorials/10_async/00_base/030_tracing/README.md

@@ -52,7 +52,9 @@ Token usage on spans is what the backend bills from, and it reads two shapes:
   (litellm, OpenAI Agents SDK, LangGraph) emit this automatically. Summed only when no
   aggregate exists in the trace.
 
-Record the turn rollup with `adk.tracing.turn_span()` instead of hand-writing usage keys:
+Record the turn rollup with `adk.tracing.turn_span()` instead of hand-writing usage keys.
+It accepts the harness `TurnUsage` that every turn adapter reports (`LangGraphTurn.usage()`,
+`run_turn(...).usage`, `ClaudeCodeTurn.usage()`, ...), cost included:
 
 ```python
 async with adk.tracing.turn_span(
@@ -61,9 +63,9 @@ async with adk.tracing.turn_span(
     input={"prompt": prompt},
     task_id=task.id,
 ) as turn:
-    result = await run_llm_calls()
-    turn.output = {"response": result.text}
-    turn.record_usage(usage=result.usage, cost_usd=result.cost_usd)
+    result = await run_turn(...)
+    turn.output = {"response": result.final_output}
+    turn.record_usage(result.usage)  # TurnUsage; cost_usd stamped automatically
 ```
 
 **Never put usage on both a rollup span's `output` and its per-call children's

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

@@ -20,7 +20,7 @@
     TracingActivityName,
 )
 from agentex.lib.core.tracing.tracer import AsyncTracer
-from agentex.lib.core.tracing.usage import usage_from_counts, validate_usage_blob
+from agentex.lib.core.harness.types import TurnUsage
 from agentex.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.model_utils import BaseModel
@@ -31,6 +31,21 @@
 DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
 TEMPORAL_SPAN_ACTIVITY_DROPPED_METRIC = "agentex.tracing.temporal_span_activity.dropped"
 
+# Token key spellings the backend accepts when billing usage from spans.
+RECOGNIZED_USAGE_KEYS = frozenset(
+    {
+        "input_tokens",
+        "prompt_tokens",
+        "output_tokens",
+        "completion_tokens",
+        "cached_input_tokens",
+        "cached_tokens",
+        "reasoning_tokens",
+        "total_tokens",
+        "cost_usd",
+    }
+)
+
 
 def _record_temporal_span_activity_dropped(event_type: str) -> None:
     try:
@@ -62,36 +77,38 @@ def __init__(self, span: Span | None):
 
     def record_usage(
         self,
-        usage: dict[str, Any] | None = None,
+        usage: TurnUsage | dict[str, Any] | None = None,
         cost_usd: float | None = None,
-        *,
-        input_tokens: int | None = None,
-        output_tokens: int | None = None,
-        total_tokens: int | None = None,
-        cached_input_tokens: int | None = None,
-        reasoning_tokens: int | None = None,
     ) -> None:
         """Record the turn's aggregate usage on the span's ``data``.
 
-        Pass either a prebuilt ``usage`` mapping (framework token spellings
-        like ``prompt_tokens``/``completion_tokens`` are accepted by the
-        backend) or individual token counts, plus an optional ``cost_usd``.
-        Individual counts are merged over the ``usage`` mapping. The usage must
-        be this turn's own tokens, not a session-cumulative total.
+        Pass the harness ``TurnUsage`` (e.g. ``LangGraphTurn.usage()`` or
+        ``run_turn(...).usage``) — its ``cost_usd`` is stamped automatically —
+        or a plain dict with backend-recognized token spellings
+        (``prompt_tokens``/``completion_tokens`` also work). An explicit
+        ``cost_usd`` argument overrides any cost carried by ``usage``. The
+        usage must be this turn's own tokens, not a session-cumulative total.
         """
         if self.span is None:
             return
 
-        blob: dict[str, Any] = validate_usage_blob(usage) if usage else {}
-        blob.update(
-            usage_from_counts(
-                input_tokens=input_tokens,
-                output_tokens=output_tokens,
-                total_tokens=total_tokens,
-                cached_input_tokens=cached_input_tokens,
-                reasoning_tokens=reasoning_tokens,
-            )
-        )
+        blob: dict[str, Any]
+        if isinstance(usage, TurnUsage):
+            blob = usage.model_dump(exclude_none=True)
+            # cost lives beside the blob as data["cost_usd"], not inside it
+            blob_cost = blob.pop("cost_usd", None)
+            if cost_usd is None:
+                cost_usd = blob_cost
+        elif usage is not None:
+            blob = dict(usage)
+            if not any(key in RECOGNIZED_USAGE_KEYS for key in blob):
+                logger.warning(
+                    "TurnSpan.record_usage: usage has no recognized token keys and will "
+                    f"not be billed. Got keys {sorted(blob)}; expected any of "
+                    f"{sorted(RECOGNIZED_USAGE_KEYS)}."
+                )
+        else:
+            blob = {}
 
         if self.span.data is not None and not isinstance(self.span.data, dict):
             logger.warning(
@@ -250,12 +267,12 @@ async def turn_span(
         Per-call child spans (LLM adapters) may still carry
         ``output["usage"]``; the backend de-dups them against this aggregate.
 
-        Example::
+        Example (with a harness turn, e.g. ``LangGraphTurn`` / ``run_turn``)::
 
             async with adk.tracing.turn_span(trace_id=task.id, name="turn", input={...}, task_id=task.id) as turn:
-                result = await run_llm_calls()
-                turn.output = {"response": result.text}
-                turn.record_usage(usage=result.usage, cost_usd=result.cost_usd)
+                result = await run_turn(...)
+                turn.output = {"response": result.final_output}
+                turn.record_usage(result.usage)  # TurnUsage, cost_usd included
         """
         async with self.span(
             trace_id=trace_id,

src/agentex/lib/core/temporal/plugins/__init__.py

@@ -52,4 +52,4 @@
     "streaming_parent_span_id",
     "TemporalStreamingHooks",
     "stream_lifecycle_content",
-]
+]

src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py

@@ -12,4 +12,4 @@
 __all__ = [
     "TemporalStreamingModel",
     "TemporalStreamingModelProvider",
-]
+]

src/agentex/lib/core/tracing/usage.py

@@ -8,6 +8,7 @@
 
 import agentex.lib.adk._modules.tracing as _tracing_mod
 from agentex.types.span import Span
+from agentex.lib.core.harness.types import TurnUsage
 from agentex.lib.adk._modules.tracing import TurnSpan, TracingModule
 from agentex.lib.core.services.adk.tracing import TracingService
 
@@ -289,23 +290,63 @@ async def test_turn_span_records_aggregate_usage_in_data(self):
         # The aggregate lives in data, never in output — output stays payload-only
         assert ended_span.output == {"response": "hello"}
 
-    async def test_turn_span_record_usage_with_individual_counts(self):
+    async def test_turn_span_record_usage_with_turn_usage(self):
         mock_service, module = _make_module()
         started = _make_span()
         mock_service.start_span.return_value = started
         mock_service.end_span.return_value = started
 
+        turn_usage = TurnUsage(
+            model="gpt-4o",
+            input_tokens=10,
+            output_tokens=5,
+            cached_input_tokens=2,
+            total_tokens=15,
+            cost_usd=0.5,
+        )
         with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
             async with module.turn_span(trace_id="trace-123", name="turn") as turn:
-                turn.record_usage(input_tokens=10, output_tokens=5, cached_input_tokens=2)
+                turn.record_usage(turn_usage)
 
         ended_span = mock_service.end_span.call_args.kwargs["span"]
+        # cost_usd is lifted out of the blob to data["cost_usd"]
+        assert ended_span.data["cost_usd"] == 0.5
         assert ended_span.data["usage"] == {
+            "model": "gpt-4o",
             "input_tokens": 10,
             "output_tokens": 5,
             "cached_input_tokens": 2,
+            "total_tokens": 15,
+            "num_tool_calls": 0,
+            "num_reasoning_blocks": 0,
         }
 
+    async def test_turn_span_explicit_cost_overrides_turn_usage_cost(self):
+        mock_service, module = _make_module()
+        started = _make_span()
+        mock_service.start_span.return_value = started
+        mock_service.end_span.return_value = started
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
+            async with module.turn_span(trace_id="trace-123", name="turn") as turn:
+                turn.record_usage(TurnUsage(input_tokens=1, cost_usd=0.5), cost_usd=0.75)
+
+        ended_span = mock_service.end_span.call_args.kwargs["span"]
+        assert ended_span.data["cost_usd"] == 0.75
+
+    async def test_turn_span_warns_on_unrecognized_usage_keys(self, caplog):
+        mock_service, module = _make_module()
+        started = _make_span()
+        mock_service.start_span.return_value = started
+        mock_service.end_span.return_value = started
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
+            async with module.turn_span(trace_id="trace-123", name="turn") as turn:
+                with caplog.at_level("WARNING"):
+                    turn.record_usage(usage={"inputTokens": 10})
+
+        assert any("no recognized token keys" in message for message in caplog.messages)
+
     async def test_turn_span_preserves_existing_data(self):
         mock_service, module = _make_module()
         started = _make_span(data={"custom": "value"})