feat(tracing): emit token usage on spans for SGP billing

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

@@ -41,6 +41,39 @@ 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.
+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(
+    trace_id=task.id,
+    name="turn",
+    input={"prompt": prompt},
+    task_id=task.id,
+) as turn:
+    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
+`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

@@ -27,7 +27,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
 
 # Unified harness surface (AGX1-375)
 from agentex.lib.core.harness import (
@@ -66,6 +66,7 @@
     "tracing",
     "events",
     "agent_task_tracker",
+    "TurnSpan",
     # Checkpointing / LangGraph
     "create_checkpointer",
     "stream_langgraph_events",

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

@@ -20,6 +20,7 @@
     TracingActivityName,
 )
 from agentex.lib.core.tracing.tracer import AsyncTracer
+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
@@ -30,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:
@@ -42,6 +58,80 @@ def _record_temporal_span_activity_dropped(event_type: str) -> None:
         pass
 
 
+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: TurnUsage | dict[str, Any] | None = None,
+        cost_usd: float | None = None,
+    ) -> None:
+        """Record the turn's aggregate usage on the span's ``data``.
+
+        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]
+        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(
+                f"TurnSpan.record_usage: span.data is {type(self.span.data).__name__} "
+                "(expected dict or None); existing data will be replaced."
+            )
+        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.
@@ -156,6 +246,47 @@ 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 (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_turn(...)
+                turn.output = {"response": result.final_output}
+                turn.record_usage(result.usage)  # TurnUsage, cost_usd included
+        """
+        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,

src/agentex/lib/core/services/adk/providers/litellm.py

@@ -25,6 +25,20 @@
 logger = logging.make_logger(__name__)
 
 
+def _stream_kwargs_with_usage(llm_config: LLMConfig) -> dict:
+    """Completion kwargs with usage reporting enabled on the final stream chunk.
+
+    litellm only reports usage for streaming calls when
+    ``stream_options={"include_usage": True}`` is set; default it on so usage
+    reaches the span. Callers can still opt out by explicitly passing
+    ``stream_options={"include_usage": False}``.
+    """
+    kwargs = llm_config.model_dump()
+    stream_options = kwargs.get("stream_options") or {}
+    kwargs["stream_options"] = {"include_usage": True, **stream_options}
+    return kwargs
+
+
 class LiteLLMService:
     def __init__(
         self,
@@ -121,7 +135,11 @@ async def chat_completion_auto_send(
 
             if span:
                 if streaming_context.task_message:
-                    span.output = streaming_context.task_message.model_dump()
+                    output = streaming_context.task_message.model_dump()
+                    # Per-call usage for billing; deduped against any turn aggregate
+                    if completion.usage is not None:
+                        output["usage"] = completion.usage.model_dump()
+                    span.output = output
             return streaming_context.task_message if streaming_context.task_message else None
 
     async def chat_completion_stream(
@@ -150,18 +168,21 @@ async def chat_completion_stream(
         if self.llm_gateway is None:
             raise ValueError("LLM Gateway is not set")
 
+        completion_kwargs = _stream_kwargs_with_usage(llm_config)
         trace = self.tracer.trace(trace_id)
         async with trace.span(
             parent_id=parent_span_id,
             name="chat_completion_stream",
-            input=llm_config.model_dump(),
+            input=completion_kwargs,
         ) as span:
             # Direct streaming outside temporal - yield each chunk as it comes
             chunks: list[Completion] = []
-            async for chunk in self.llm_gateway.acompletion_stream(**llm_config.model_dump()):
+            async for chunk in self.llm_gateway.acompletion_stream(**completion_kwargs):
                 chunks.append(chunk)
                 yield chunk
             if span:
+                # The usage-bearing final chunk survives concat, so the dumped
+                # completion carries usage for billing
                 span.output = concat_completion_chunks(chunks).model_dump()
 
     async def chat_completion_stream_auto_send(
@@ -190,11 +211,12 @@ async def chat_completion_stream_auto_send(
         if not llm_config.stream:
             llm_config.stream = True
 
+        completion_kwargs = _stream_kwargs_with_usage(llm_config)
         trace = self.tracer.trace(trace_id)
         async with trace.span(
             parent_id=parent_span_id,
             name="chat_completion_stream_auto_send",
-            input=llm_config.model_dump(),
+            input=completion_kwargs,
         ) as span:
             # Use streaming context manager
             async with self.streaming_service.streaming_task_message_context(
@@ -208,8 +230,11 @@ async def chat_completion_stream_auto_send(
             ) as streaming_context:
                 # Get the streaming response
                 chunks = []
-                async for response in self.llm_gateway.acompletion_stream(**llm_config.model_dump()):
+                async for response in self.llm_gateway.acompletion_stream(**completion_kwargs):
                     heartbeat_if_in_workflow("chat completion streaming")
+                    # Store every chunk for final message assembly, including
+                    # the usage-only final chunk, which has no choices
+                    chunks.append(response)
                     if response.choices and len(response.choices) > 0 and response.choices[0].delta:
                         delta = response.choices[0].delta.content
                         if delta:
@@ -223,9 +248,6 @@ async def chat_completion_stream_auto_send(
                             )
                             heartbeat_if_in_workflow("content chunk streamed")
 
-                        # Store the chunk for final message assembly
-                        chunks.append(response)
-
                 # Update the final message content
                 complete_message = concat_completion_chunks(chunks)
                 if complete_message and complete_message.choices and complete_message.choices[0].message:
@@ -246,6 +268,10 @@ async def chat_completion_stream_auto_send(
 
             if span:
                 if streaming_context.task_message:
-                    span.output = streaming_context.task_message.model_dump()
+                    output = streaming_context.task_message.model_dump()
+                    # Per-call usage for billing; deduped against any turn aggregate
+                    if complete_message.usage is not None:
+                        output["usage"] = complete_message.usage.model_dump()
+                    span.output = output
 
         return streaming_context.task_message if streaming_context.task_message else None

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

@@ -11,7 +11,6 @@
 Example:
     >>> from agentex.lib.core.temporal.plugins.openai_agents import (
     ...     TemporalStreamingModelProvider,
-    ...     TemporalTracingModelProvider,
     ...     ContextInterceptor,
     ... )
     >>> from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters
@@ -37,7 +36,6 @@
     ContextInterceptor,
     TemporalStreamingHooks,
     TemporalStreamingModel,
-    TemporalTracingModelProvider,
     TemporalStreamingModelProvider,
     streaming_task_id,
     streaming_trace_id,
@@ -48,11 +46,10 @@
 __all__ = [
     "TemporalStreamingModel",
     "TemporalStreamingModelProvider",
-    "TemporalTracingModelProvider",
     "ContextInterceptor",
     "streaming_task_id",
     "streaming_trace_id",
     "streaming_parent_span_id",
     "TemporalStreamingHooks",
     "stream_lifecycle_content",
-]
+]

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

@@ -61,9 +61,6 @@
 from agentex.lib.core.temporal.plugins.openai_agents.hooks.activities import (
     stream_lifecycle_content,
 )
-from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
-    TemporalTracingModelProvider,
-)
 from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model import (
     TemporalStreamingModel,
     TemporalStreamingModelProvider,
@@ -78,7 +75,6 @@
 __all__ = [
     "TemporalStreamingModel",
     "TemporalStreamingModelProvider",
-    "TemporalTracingModelProvider",
     "ContextInterceptor",
     "streaming_task_id",
     "streaming_trace_id",

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

@@ -4,11 +4,6 @@
 capabilities to standard OpenAI models when running in Temporal workflows/activities.
 """
 
-from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
-    TemporalTracingModelProvider,
-    TemporalTracingResponsesModel,
-    TemporalTracingChatCompletionsModel,
-)
 from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model import (
     TemporalStreamingModel,
     TemporalStreamingModelProvider,
@@ -17,7 +12,4 @@
 __all__ = [
     "TemporalStreamingModel",
     "TemporalStreamingModelProvider",
-    "TemporalTracingModelProvider",
-    "TemporalTracingResponsesModel",
-    "TemporalTracingChatCompletionsModel",
-]
+]