feat(adk): add turn_span helper for per-turn aggregate usage and cost

TurnSpan.record_usage writes the billable aggregate to span.data (usage + cost_usd), encapsulating the contract so agents cannot re-introduce the double-count bug. Documents the usage/cost span contract in the tracing tutorial.

Author: levilentz

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

@@ -41,6 +41,37 @@ await adk.tracing.end_span(
 
 Spans create a hierarchical view of agent execution, making it easy to see which operations take time and where errors occur.
 
+## Token Usage & Cost Tracking
+
+Token usage on spans is what the backend bills from, and it reads two shapes:
+
+- **Per-turn aggregate** — `span.data["usage"]` + `span.data["cost_usd"]`. Emit at most
+  once per turn, holding that turn's own usage (not a session-cumulative total). When a
+  trace has an aggregate, the backend keeps it and de-dups all per-call spans against it.
+- **Per-call detail** — `span.output["usage"]`. Optional; the SDK's LLM adapters
+  (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:
+
+```python
+async with adk.tracing.turn_span(
+    trace_id=task.id,
+    name="turn",
+    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)
+```
+
+**Never put usage on both a rollup span's `output` and its per-call children's
+`output`** — that double-counts. `turn_span` writes the aggregate to `data`, so child
+spans stay safe to emit. Recognized token keys: `input_tokens`/`prompt_tokens`,
+`output_tokens`/`completion_tokens`, `cached_input_tokens`/`cached_tokens`,
+`reasoning_tokens`; cost is `cost_usd`.
+
 ## When to Use
 - Debugging complex agent behaviors
 - Performance optimization and bottleneck identification

src/agentex/lib/adk/__init__.py

@@ -14,7 +14,7 @@
 from agentex.lib.adk._modules.state import StateModule
 from agentex.lib.adk._modules.streaming import StreamingModule
 from agentex.lib.adk._modules.tasks import TasksModule
-from agentex.lib.adk._modules.tracing import TracingModule
+from agentex.lib.adk._modules.tracing import TracingModule, TurnSpan
 
 from agentex.lib.adk import providers
 from agentex.lib.adk import utils
@@ -40,6 +40,7 @@
     "tracing",
     "events",
     "agent_task_tracker",
+    "TurnSpan",
 
     # Checkpointing / LangGraph
     "create_checkpointer",

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

@@ -18,6 +18,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.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.model_utils import BaseModel
@@ -28,6 +29,73 @@
 DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
 
 
+class TurnSpan:
+    """Handle for a turn-level (rollup) span, yielded by ``TracingModule.turn_span``.
+
+    Encapsulates the billing contract so agents cannot double-count usage:
+    the turn's aggregate usage goes to ``span.data["usage"]`` (+
+    ``span.data["cost_usd"]``) via :meth:`record_usage`. The backend keeps the
+    aggregate and de-dups any per-call ``output["usage"]`` children against it.
+    Never hand-write usage into ``output`` on a rollup span — that is the
+    double-count bug this helper exists to prevent.
+
+    All methods no-op when tracing is disabled (``span`` is None), so agent
+    code needs no ``if span:`` guards.
+    """
+
+    def __init__(self, span: Span | None):
+        self.span = span
+
+    def record_usage(
+        self,
+        usage: 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.
+        """
+        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,
+            )
+        )
+
+        data = self.span.data if isinstance(self.span.data, dict) else {}
+        if blob:
+            data["usage"] = blob
+        if cost_usd is not None:
+            data["cost_usd"] = cost_usd
+        self.span.data = data
+
+    @property
+    def output(self) -> Any:
+        return self.span.output if self.span is not None else None
+
+    @output.setter
+    def output(self, value: Any) -> None:
+        if self.span is not None:
+            self.span.output = value
+
+
 class TracingModule:
     """
     Module for managing tracing and span operations in Agentex.
@@ -143,6 +211,49 @@ async def span(
                     retry_policy=retry_policy,
                 )
 
+    @asynccontextmanager
+    async def turn_span(
+        self,
+        trace_id: str,
+        name: str,
+        input: list[Any] | dict[str, Any] | BaseModel | None = None,
+        data: list[Any] | dict[str, Any] | BaseModel | None = None,
+        parent_id: str | None = None,
+        task_id: str | None = None,
+        start_to_close_timeout: timedelta = timedelta(seconds=5),
+        heartbeat_timeout: timedelta = timedelta(seconds=5),
+        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
+    ) -> AsyncGenerator[TurnSpan, None]:
+        """Span for one agent turn, with usage recorded as the billable aggregate.
+
+        Same lifecycle as :meth:`span`, but yields a :class:`TurnSpan` whose
+        ``record_usage(usage=..., cost_usd=...)`` writes the turn's rollup
+        usage to ``span.data`` — the shape the backend bills once per turn.
+        Per-call child spans (LLM adapters) may still carry
+        ``output["usage"]``; the backend de-dups them against this aggregate.
+
+        Example::
+
+            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)
+        """
+        async with self.span(
+            trace_id=trace_id,
+            name=name,
+            input=input,
+            data=data,
+            parent_id=parent_id,
+            task_id=task_id,
+            start_to_close_timeout=start_to_close_timeout,
+            heartbeat_timeout=heartbeat_timeout,
+            retry_policy=retry_policy,
+        ) as span:
+            yield TurnSpan(span)
+
     async def start_span(
         self,
         trace_id: str,

tests/lib/adk/test_tracing_module.py

@@ -5,7 +5,7 @@
 
 import agentex.lib.adk._modules.tracing as _tracing_mod
 from agentex.types.span import Span
-from agentex.lib.adk._modules.tracing import TracingModule
+from agentex.lib.adk._modules.tracing import TurnSpan, TracingModule
 from agentex.lib.core.services.adk.tracing import TracingService
 
 
@@ -115,3 +115,95 @@ async def test_span_context_manager_noop_when_no_trace_id(self):
 
         mock_service.start_span.assert_not_called()
         mock_service.end_span.assert_not_called()
+
+
+class TestTurnSpan:
+    async def test_turn_span_records_aggregate_usage_in_data(self):
+        mock_service, module = _make_module()
+        started = _make_span(task_id="task-abc")
+        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",
+                task_id="task-abc",
+            ) as turn:
+                assert isinstance(turn, TurnSpan)
+                turn.output = {"response": "hello"}
+                turn.record_usage(
+                    usage={"input_tokens": 100, "output_tokens": 40, "total_tokens": 140},
+                    cost_usd=0.0125,
+                )
+
+        ended_span = mock_service.end_span.call_args.kwargs["span"]
+        assert ended_span.data["usage"] == {
+            "input_tokens": 100,
+            "output_tokens": 40,
+            "total_tokens": 140,
+        }
+        assert ended_span.data["cost_usd"] == 0.0125
+        # 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):
+        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(input_tokens=10, output_tokens=5, cached_input_tokens=2)
+
+        ended_span = mock_service.end_span.call_args.kwargs["span"]
+        assert ended_span.data["usage"] == {
+            "input_tokens": 10,
+            "output_tokens": 5,
+            "cached_input_tokens": 2,
+        }
+
+    async def test_turn_span_preserves_existing_data(self):
+        mock_service, module = _make_module()
+        started = _make_span(data={"custom": "value"})
+        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", data={"custom": "value"}
+            ) as turn:
+                turn.record_usage(usage={"prompt_tokens": 3, "completion_tokens": 4})
+
+        ended_span = mock_service.end_span.call_args.kwargs["span"]
+        assert ended_span.data["custom"] == "value"
+        assert ended_span.data["usage"] == {"prompt_tokens": 3, "completion_tokens": 4}
+
+    async def test_turn_span_cost_only(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(cost_usd=0.5)
+
+        ended_span = mock_service.end_span.call_args.kwargs["span"]
+        assert ended_span.data == {"cost_usd": 0.5}
+        assert "usage" not in ended_span.data
+
+    async def test_turn_span_noop_when_no_trace_id(self):
+        mock_service, module = _make_module()
+
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=False):
+            async with module.turn_span(trace_id="", name="turn") as turn:
+                assert turn.span is None
+                # Must not raise when tracing is disabled
+                turn.record_usage(usage={"input_tokens": 1}, cost_usd=0.1)
+                turn.output = {"response": "x"}
+                assert turn.output is None
+
+        mock_service.start_span.assert_not_called()
+        mock_service.end_span.assert_not_called()