From d3ca011cb8c081c6fb5e87eb911dd67eb0d59687 Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 13:28:52 -0700
Subject: [PATCH 1/9] feat(tracing): add usage normalization helpers for span
 billing contract

---
 src/agentex/lib/core/tracing/usage.py | 115 ++++++++++++++++++++++++++
 tests/lib/core/tracing/test_usage.py  |  90 ++++++++++++++++++++
 2 files changed, 205 insertions(+)
 create mode 100644 src/agentex/lib/core/tracing/usage.py
 create mode 100644 tests/lib/core/tracing/test_usage.py

diff --git a/src/agentex/lib/core/tracing/usage.py b/src/agentex/lib/core/tracing/usage.py
new file mode 100644
index 000000000..63dbfcd80
--- /dev/null
+++ b/src/agentex/lib/core/tracing/usage.py
@@ -0,0 +1,115 @@
+"""Helpers for putting LLM token usage onto trace spans in the billable shape.
+
+The AgentEx backend bills token usage from ``application_trace_span`` rows using
+two span fields:
+
+- ``span.data["usage"]`` (+ ``span.data["cost_usd"]``): the per-turn AGGREGATE.
+  Emit at most once per turn, holding that turn's own (per-invocation, not
+  session-cumulative) usage. When a trace contains an aggregate, the backend
+  keeps it and drops all per-call spans in that trace.
+- ``span.output["usage"]``: per-call detail. Summed by the backend, and dropped
+  whenever an aggregate exists in the trace.
+
+Never emit usage on both a rollup span and its per-call children via
+``output["usage"]`` — that double-counts.
+
+Recognized token key spellings (either spelling of a pair works):
+``input_tokens``/``prompt_tokens``, ``output_tokens``/``completion_tokens``,
+``cached_input_tokens``/``cached_tokens``, ``reasoning_tokens``; cost is
+``cost_usd``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from collections.abc import Mapping
+
+from agentex.lib.utils.logging import make_logger
+
+logger = make_logger(__name__)
+
+# Key spellings the backend accepts when summing token 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 usage_from_counts(
+    *,
+    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,
+    cost_usd: float | None = None,
+) -> dict[str, Any]:
+    """Build a usage blob with canonical key spellings, omitting None values."""
+    usage: dict[str, Any] = {}
+    if input_tokens is not None:
+        usage["input_tokens"] = input_tokens
+    if output_tokens is not None:
+        usage["output_tokens"] = output_tokens
+    if total_tokens is not None:
+        usage["total_tokens"] = total_tokens
+    if cached_input_tokens is not None:
+        usage["cached_input_tokens"] = cached_input_tokens
+    if reasoning_tokens is not None:
+        usage["reasoning_tokens"] = reasoning_tokens
+    if cost_usd is not None:
+        usage["cost_usd"] = cost_usd
+    return usage
+
+
+def usage_from_openai_response_usage(usage: Any) -> dict[str, Any] | None:
+    """Extract a usage blob from an OpenAI-style usage object.
+
+    Duck-typed so it works for both ``agents.usage.Usage`` (OpenAI Agents SDK)
+    and ``openai.types.responses.ResponseUsage``: reads ``input_tokens``,
+    ``output_tokens``, ``total_tokens``, ``input_tokens_details.cached_tokens``,
+    and ``output_tokens_details.reasoning_tokens``.
+
+    Returns None when there is nothing usable to report.
+    """
+    if usage is None:
+        return None
+
+    input_tokens = getattr(usage, "input_tokens", None)
+    output_tokens = getattr(usage, "output_tokens", None)
+    if input_tokens is None and output_tokens is None:
+        return None
+
+    input_details = getattr(usage, "input_tokens_details", None)
+    output_details = getattr(usage, "output_tokens_details", None)
+    return usage_from_counts(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=getattr(usage, "total_tokens", None),
+        cached_input_tokens=getattr(input_details, "cached_tokens", None) if input_details is not None else None,
+        reasoning_tokens=getattr(output_details, "reasoning_tokens", None) if output_details is not None else None,
+    )
+
+
+def validate_usage_blob(usage: Mapping[str, Any]) -> dict[str, Any]:
+    """Return the usage mapping as a plain dict, warning on unrecognized shapes.
+
+    The blob is passed through untouched so callers keep full control; the
+    warning catches typos (e.g. ``inputTokens``) that the backend would
+    silently ignore when billing.
+    """
+    blob = dict(usage)
+    if not any(key in RECOGNIZED_USAGE_KEYS for key in blob):
+        logger.warning(
+            "Usage blob has no recognized token keys and will not be billed. "
+            f"Got keys {sorted(blob)}; expected any of {sorted(RECOGNIZED_USAGE_KEYS)}."
+        )
+    return blob
diff --git a/tests/lib/core/tracing/test_usage.py b/tests/lib/core/tracing/test_usage.py
new file mode 100644
index 000000000..110cc07e3
--- /dev/null
+++ b/tests/lib/core/tracing/test_usage.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+from agentex.lib.core.tracing.usage import (
+    usage_from_counts,
+    validate_usage_blob,
+    usage_from_openai_response_usage,
+)
+
+
+class TestUsageFromCounts:
+    def test_all_fields(self):
+        assert usage_from_counts(
+            input_tokens=100,
+            output_tokens=50,
+            total_tokens=150,
+            cached_input_tokens=20,
+            reasoning_tokens=10,
+            cost_usd=0.0123,
+        ) == {
+            "input_tokens": 100,
+            "output_tokens": 50,
+            "total_tokens": 150,
+            "cached_input_tokens": 20,
+            "reasoning_tokens": 10,
+            "cost_usd": 0.0123,
+        }
+
+    def test_none_values_omitted(self):
+        assert usage_from_counts(input_tokens=100, output_tokens=50) == {
+            "input_tokens": 100,
+            "output_tokens": 50,
+        }
+
+    def test_explicit_zeros_kept(self):
+        assert usage_from_counts(input_tokens=0, output_tokens=0) == {
+            "input_tokens": 0,
+            "output_tokens": 0,
+        }
+
+
+class TestUsageFromOpenAIResponseUsage:
+    def test_none_returns_none(self):
+        assert usage_from_openai_response_usage(None) is None
+
+    def test_object_without_token_fields_returns_none(self):
+        assert usage_from_openai_response_usage(SimpleNamespace(requests=1)) is None
+
+    def test_full_usage_with_details(self):
+        usage = SimpleNamespace(
+            input_tokens=120,
+            output_tokens=80,
+            total_tokens=200,
+            input_tokens_details=SimpleNamespace(cached_tokens=30),
+            output_tokens_details=SimpleNamespace(reasoning_tokens=40),
+        )
+        assert usage_from_openai_response_usage(usage) == {
+            "input_tokens": 120,
+            "output_tokens": 80,
+            "total_tokens": 200,
+            "cached_input_tokens": 30,
+            "reasoning_tokens": 40,
+        }
+
+    def test_usage_without_details(self):
+        usage = SimpleNamespace(
+            input_tokens=10,
+            output_tokens=5,
+            total_tokens=15,
+            input_tokens_details=None,
+            output_tokens_details=None,
+        )
+        assert usage_from_openai_response_usage(usage) == {
+            "input_tokens": 10,
+            "output_tokens": 5,
+            "total_tokens": 15,
+        }
+
+
+class TestValidateUsageBlob:
+    def test_passthrough_recognized_keys(self):
+        blob = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert validate_usage_blob(blob) == blob
+
+    def test_warns_on_unrecognized_keys(self, caplog):
+        with caplog.at_level("WARNING"):
+            result = validate_usage_blob({"inputTokens": 10})
+        assert result == {"inputTokens": 10}
+        assert any("no recognized token keys" in message for message in caplog.messages)

From c0aba7681da1f8ad3919a2dd32bca636855cbaa8 Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 13:32:25 -0700
Subject: [PATCH 2/9] feat(tracing): emit real token usage from openai_agents
 temporal models

Captures ResponseCompletedEvent usage in the streaming model (was zeroed) and response.usage in both tracing wrappers, writing span.output.usage for billing. Also implements stream_response on the tracing wrappers, which were abstract and raised TypeError on instantiation.
---
 .../models/temporal_tracing_model.py          |  29 ++-
 tests/lib/core/temporal/plugins/__init__.py   |   0
 .../plugins/openai_agents/__init__.py         |   0
 .../plugins/openai_agents/test_model_usage.py | 238 ++++++++++++++++++
 4 files changed, 264 insertions(+), 3 deletions(-)
 create mode 100644 tests/lib/core/temporal/plugins/__init__.py
 create mode 100644 tests/lib/core/temporal/plugins/openai_agents/__init__.py
 create mode 100644 tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py

diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
index e3d1b3804..02372aedd 100644
--- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
+++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
@@ -28,6 +28,7 @@
 from agents.models.openai_responses import OpenAIResponsesModel
 from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
 
+from agentex.lib.core.tracing.usage import usage_from_openai_response_usage
 from agentex.lib.core.tracing.tracer import AsyncTracer
 
 # Import AgentEx components
@@ -243,10 +244,15 @@ async def get_response(
                                 continue
 
                     # Set span output with structured data
-                    span.output = {  # type: ignore[attr-defined]
+                    output_data: dict[str, Any] = {
                         "new_items": new_items,
                         "final_output": final_output,
                     }
+                    # Per-call usage for billing; deduped against any turn aggregate
+                    usage_blob = usage_from_openai_response_usage(getattr(response, "usage", None))
+                    if usage_blob:
+                        output_data["usage"] = usage_blob
+                    span.output = output_data  # type: ignore[attr-defined]
 
                     return response
 
@@ -271,6 +277,12 @@ async def get_response(
                 **kwargs,
             )
 
+    @override
+    def stream_response(self, *args, **kwargs):
+        """Streaming is handled via get_response in Temporal activities.
+        Required so the class is concrete and instantiable at runtime."""
+        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
+
 
 class TemporalTracingChatCompletionsModel(Model):
     """Wrapper for OpenAIChatCompletionsModel that adds AgentEx tracing.
@@ -376,10 +388,15 @@ async def get_response(
                                 continue
 
                     # Set span output with structured data
-                    span.output = {  # type: ignore[attr-defined]
+                    output_data: dict[str, Any] = {
                         "new_items": new_items,
                         "final_output": final_output,
                     }
+                    # Per-call usage for billing; deduped against any turn aggregate
+                    usage_blob = usage_from_openai_response_usage(getattr(response, "usage", None))
+                    if usage_blob:
+                        output_data["usage"] = usage_blob
+                    span.output = output_data  # type: ignore[attr-defined]
 
                     return response
 
@@ -399,4 +416,10 @@ async def get_response(
                 handoffs=handoffs,
                 tracing=tracing,
                 **kwargs,
-            )
\ No newline at end of file
+            )
+
+    @override
+    def stream_response(self, *args, **kwargs):
+        """Streaming is handled via get_response in Temporal activities.
+        Required so the class is concrete and instantiable at runtime."""
+        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
\ No newline at end of file
diff --git a/tests/lib/core/temporal/plugins/__init__.py b/tests/lib/core/temporal/plugins/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/lib/core/temporal/plugins/openai_agents/__init__.py b/tests/lib/core/temporal/plugins/openai_agents/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
new file mode 100644
index 000000000..26db6fb5f
--- /dev/null
+++ b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
@@ -0,0 +1,238 @@
+"""Tests that the openai_agents temporal models copy real token usage onto spans.
+
+The backend bills per-call usage from ``span.output["usage"]``; these tests
+assert each model writes the framework-reported usage there (and, for the
+streaming model, into the returned ``ModelResponse.usage``) instead of
+dropping it.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from agents import ModelResponse
+from agents.usage import Usage, InputTokensDetails, OutputTokensDetails
+
+from agentex.types.span import Span
+
+pytestmark = pytest.mark.asyncio
+
+
+class FakeTrace:
+    """Captures spans handed out by trace.span() so tests can inspect them."""
+
+    def __init__(self) -> None:
+        self.spans: list[Span] = []
+
+    @asynccontextmanager
+    async def span(self, name, parent_id=None, input=None, data=None, task_id=None):
+        span = Span(
+            id=f"span-{len(self.spans)}",
+            name=name,
+            start_time=datetime.now(UTC),
+            trace_id="trace-1",
+            parent_id=parent_id,
+            input=input,
+            data=data,
+            task_id=task_id,
+        )
+        self.spans.append(span)
+        yield span
+
+
+class FakeTracer:
+    def __init__(self) -> None:
+        self.trace_obj = FakeTrace()
+
+    def trace(self, trace_id):
+        return self.trace_obj
+
+
+@pytest.fixture
+def tracing_contextvars():
+    from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
+        streaming_task_id,
+        streaming_trace_id,
+        streaming_parent_span_id,
+    )
+
+    tokens = [
+        streaming_task_id.set("task-1"),
+        streaming_trace_id.set("trace-1"),
+        streaming_parent_span_id.set("parent-span-1"),
+    ]
+    yield
+    streaming_task_id.reset(tokens[0])
+    streaming_trace_id.reset(tokens[1])
+    streaming_parent_span_id.reset(tokens[2])
+
+
+def _agents_usage() -> Usage:
+    return Usage(
+        requests=1,
+        input_tokens=120,
+        output_tokens=80,
+        total_tokens=200,
+        input_tokens_details=InputTokensDetails(cached_tokens=30),
+        output_tokens_details=OutputTokensDetails(reasoning_tokens=40),
+    )
+
+EXPECTED_USAGE_BLOB = {
+    "input_tokens": 120,
+    "output_tokens": 80,
+    "total_tokens": 200,
+    "cached_input_tokens": 30,
+    "reasoning_tokens": 40,
+}
+
+
+class TestTemporalTracingModels:
+    async def _run_wrapper(self, wrapper_cls) -> Span:
+        tracer = FakeTracer()
+        base_model = MagicMock()
+        base_model.model = "gpt-4o"
+        base_model.get_response = AsyncMock(
+            return_value=ModelResponse(output=[], usage=_agents_usage(), response_id="resp-1")
+        )
+
+        model = wrapper_cls(base_model, tracer)
+        from agents import ModelSettings
+
+        response = await model.get_response(
+            system_instructions=None,
+            input="hello",
+            model_settings=ModelSettings(),
+            tools=[],
+            output_schema=None,
+            handoffs=[],
+            tracing=None,
+        )
+        assert response.usage.input_tokens == 120
+        assert len(tracer.trace_obj.spans) == 1
+        return tracer.trace_obj.spans[0]
+
+    async def test_responses_model_writes_usage_to_span_output(self, tracing_contextvars):
+        from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
+            TemporalTracingResponsesModel,
+        )
+
+        span = await self._run_wrapper(TemporalTracingResponsesModel)
+        assert span.output["usage"] == EXPECTED_USAGE_BLOB
+
+    async def test_chat_completions_model_writes_usage_to_span_output(self, tracing_contextvars):
+        from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
+            TemporalTracingChatCompletionsModel,
+        )
+
+        span = await self._run_wrapper(TemporalTracingChatCompletionsModel)
+        assert span.output["usage"] == EXPECTED_USAGE_BLOB
+
+
+class FakeStream:
+    def __init__(self, events) -> None:
+        self._events = events
+
+    def __aiter__(self):
+        async def gen():
+            for event in self._events:
+                yield event
+
+        return gen()
+
+
+class TestTemporalStreamingModel:
+    async def test_streaming_model_captures_final_response_usage(self, tracing_contextvars):
+        import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm
+
+        from openai.types.responses import Response, ResponseCompletedEvent
+        from openai.types.responses.response_usage import (
+            ResponseUsage,
+            InputTokensDetails as ResponseInputTokensDetails,
+            OutputTokensDetails as ResponseOutputTokensDetails,
+        )
+
+        usage = ResponseUsage(
+            input_tokens=120,
+            output_tokens=80,
+            total_tokens=200,
+            input_tokens_details=ResponseInputTokensDetails(cached_tokens=30),
+            output_tokens_details=ResponseOutputTokensDetails(reasoning_tokens=40),
+        )
+        completed = ResponseCompletedEvent.model_construct(
+            type="response.completed",
+            response=Response.model_construct(output=[], usage=usage),
+        )
+
+        fake_tracer = FakeTracer()
+        openai_client = MagicMock()
+        openai_client.responses.create = AsyncMock(return_value=FakeStream([completed]))
+
+        with (
+            patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()),
+            patch.object(tsm, "AsyncTracer", return_value=fake_tracer),
+        ):
+            model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client)
+
+        from agents import ModelSettings
+
+        response = await model.get_response(
+            system_instructions=None,
+            input="hello",
+            model_settings=ModelSettings(),
+            tools=[],
+            output_schema=None,
+            handoffs=[],
+            tracing=None,
+        )
+
+        # Real usage lands on the returned ModelResponse (was zeroed before)
+        assert response.usage.requests == 1
+        assert response.usage.input_tokens == 120
+        assert response.usage.output_tokens == 80
+        assert response.usage.total_tokens == 200
+        assert response.usage.input_tokens_details.cached_tokens == 30
+        assert response.usage.output_tokens_details.reasoning_tokens == 40
+
+        # And on the span output for billing
+        assert len(fake_tracer.trace_obj.spans) == 1
+        span = fake_tracer.trace_obj.spans[0]
+        assert span.output["usage"] == EXPECTED_USAGE_BLOB
+
+    async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_contextvars):
+        import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm
+
+        from openai.types.responses import Response, ResponseCompletedEvent
+
+        completed = ResponseCompletedEvent.model_construct(
+            type="response.completed",
+            response=Response.model_construct(output=[], usage=None),
+        )
+
+        fake_tracer = FakeTracer()
+        openai_client = MagicMock()
+        openai_client.responses.create = AsyncMock(return_value=FakeStream([completed]))
+
+        with (
+            patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()),
+            patch.object(tsm, "AsyncTracer", return_value=fake_tracer),
+        ):
+            model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client)
+
+        from agents import ModelSettings
+
+        response = await model.get_response(
+            system_instructions=None,
+            input="hello",
+            model_settings=ModelSettings(),
+            tools=[],
+            output_schema=None,
+            handoffs=[],
+            tracing=None,
+        )
+
+        assert response.usage.input_tokens == 0
+        span = fake_tracer.trace_obj.spans[0]
+        assert "usage" not in span.output

From 89b1b6b084e76b6c11944f501957e577db61c012 Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 13:38:07 -0700
Subject: [PATCH 3/9] feat(tracing): emit token usage from litellm streaming
 and auto_send spans

Streaming calls now default stream_options.include_usage=True, the usage-only final chunk is collected, and both auto_send variants attach completion usage to span output. concat_completion_chunks no longer drops choices when a chunk has none.
---
 .../core/services/adk/providers/litellm.py    |  44 +++-
 src/agentex/lib/utils/completions.py          |   9 +-
 tests/lib/adk/providers/test_litellm_usage.py | 213 ++++++++++++++++++
 tests/lib/utils/__init__.py                   |   0
 tests/lib/utils/test_completions.py           |  52 +++++
 5 files changed, 308 insertions(+), 10 deletions(-)
 create mode 100644 tests/lib/adk/providers/test_litellm_usage.py
 create mode 100644 tests/lib/utils/__init__.py
 create mode 100644 tests/lib/utils/test_completions.py

diff --git a/src/agentex/lib/core/services/adk/providers/litellm.py b/src/agentex/lib/core/services/adk/providers/litellm.py
index a511aa7b8..416077e8f 100644
--- a/src/agentex/lib/core/services/adk/providers/litellm.py
+++ b/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
diff --git a/src/agentex/lib/utils/completions.py b/src/agentex/lib/utils/completions.py
index fe62c7d12..595e0fb01 100644
--- a/src/agentex/lib/utils/completions.py
+++ b/src/agentex/lib/utils/completions.py
@@ -3,6 +3,7 @@
 from copy import deepcopy
 from typing import Any
 from functools import reduce, singledispatch
+from itertools import zip_longest
 
 from agentex.lib.types.llm_messages import (
     Delta,
@@ -21,7 +22,13 @@ def _concat_chunks(_a: None, b: Any):
 
 @_concat_chunks.register
 def _(a: Completion, b: Completion) -> Completion:
-    a.choices = [_concat_chunks(*c) for c in zip(a.choices, b.choices, strict=False)]
+    # Chunks can have unequal choices: with stream_options.include_usage the
+    # final chunk carries usage but no choices. Keep the unpaired side instead
+    # of truncating to the shorter list.
+    a.choices = [
+        x if y is None else y if x is None else _concat_chunks(x, y)
+        for x, y in zip_longest(a.choices, b.choices)
+    ]
     a.usage = _concat_chunks(a.usage, b.usage)
 
     return a
diff --git a/tests/lib/adk/providers/test_litellm_usage.py b/tests/lib/adk/providers/test_litellm_usage.py
new file mode 100644
index 000000000..5f3bee405
--- /dev/null
+++ b/tests/lib/adk/providers/test_litellm_usage.py
@@ -0,0 +1,213 @@
+"""Tests that LiteLLMService puts LLM token usage on spans for billing.
+
+Covers the paths that previously dropped usage: both auto_send variants (span
+output was only the TaskMessage dump) and streaming (litellm omits usage unless
+``stream_options.include_usage`` is set).
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from contextlib import asynccontextmanager
+from unittest.mock import AsyncMock, MagicMock
+
+from agentex.types.span import Span
+from agentex.types.task_message import TaskMessage
+from agentex.types.task_message_content import TextContent
+from agentex.lib.types.llm_messages import (
+    Delta,
+    Usage,
+    Choice,
+    LLMConfig,
+    Completion,
+    AssistantMessage,
+)
+from agentex.lib.core.services.adk.providers.litellm import (
+    LiteLLMService,
+    _stream_kwargs_with_usage,
+)
+
+
+class FakeTrace:
+    def __init__(self) -> None:
+        self.spans: list[Span] = []
+
+    @asynccontextmanager
+    async def span(self, name, parent_id=None, input=None, data=None, task_id=None):
+        span = Span(
+            id=f"span-{len(self.spans)}",
+            name=name,
+            start_time=datetime.now(UTC),
+            trace_id="trace-1",
+            parent_id=parent_id,
+            input=input,
+        )
+        self.spans.append(span)
+        yield span
+
+
+class FakeTracer:
+    def __init__(self) -> None:
+        self.trace_obj = FakeTrace()
+
+    def trace(self, trace_id):
+        return self.trace_obj
+
+
+def _make_streaming_service():
+    streaming_context = MagicMock()
+    streaming_context.task_message = TaskMessage(
+        id="msg-1",
+        task_id="task-1",
+        content=TextContent(author="agent", content="", format="markdown"),
+    )
+    streaming_context.stream_update = AsyncMock()
+
+    @asynccontextmanager
+    async def fake_context(**kwargs):
+        yield streaming_context
+
+    streaming_service = MagicMock()
+    streaming_service.streaming_task_message_context = fake_context
+    return streaming_service, streaming_context
+
+
+def _make_service(llm_gateway) -> tuple[LiteLLMService, FakeTracer]:
+    streaming_service, _ = _make_streaming_service()
+    tracer = FakeTracer()
+    service = LiteLLMService(
+        agentex_client=MagicMock(),
+        streaming_service=streaming_service,
+        tracer=tracer,
+        llm_gateway=llm_gateway,
+    )
+    return service, tracer
+
+
+def _stream_gateway(chunks, captured_kwargs):
+    gateway = MagicMock()
+
+    def acompletion_stream(**kwargs):
+        captured_kwargs.update(kwargs)
+
+        async def stream():
+            for chunk in chunks:
+                yield chunk
+
+        return stream()
+
+    gateway.acompletion_stream = acompletion_stream
+    return gateway
+
+
+def _delta_chunk(content: str, role: str | None = None) -> Completion:
+    return Completion(choices=[Choice(index=0, delta=Delta(content=content, role=role))])
+
+
+def _usage_only_chunk() -> Completion:
+    return Completion(
+        choices=[],
+        usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
+    )
+
+
+class TestStreamKwargsWithUsage:
+    def test_defaults_include_usage_on(self):
+        config = LLMConfig(model="gpt-4o", messages=[], stream=True)
+        assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": True}
+
+    def test_caller_opt_out_preserved(self):
+        config = LLMConfig(model="gpt-4o", messages=[], stream=True, stream_options={"include_usage": False})
+        assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": False}
+
+    def test_merges_with_other_stream_options(self):
+        config = LLMConfig(model="gpt-4o", messages=[], stream=True, stream_options={"other": 1})
+        assert _stream_kwargs_with_usage(config)["stream_options"] == {"include_usage": True, "other": 1}
+
+
+class TestChatCompletionAutoSend:
+    async def test_span_output_carries_usage(self):
+        completion = Completion(
+            choices=[Choice(index=0, message=AssistantMessage(content="Hello!"), finish_reason="stop")],
+            usage=Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10),
+        )
+        gateway = MagicMock()
+        gateway.acompletion = AsyncMock(return_value=completion)
+        service, tracer = _make_service(gateway)
+
+        await service.chat_completion_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=False),
+            trace_id="trace-1",
+        )
+
+        span = tracer.trace_obj.spans[0]
+        assert span.output["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
+
+    async def test_span_output_omits_usage_when_absent(self):
+        completion = Completion(
+            choices=[Choice(index=0, message=AssistantMessage(content="Hello!"), finish_reason="stop")],
+        )
+        gateway = MagicMock()
+        gateway.acompletion = AsyncMock(return_value=completion)
+        service, tracer = _make_service(gateway)
+
+        await service.chat_completion_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=False),
+            trace_id="trace-1",
+        )
+
+        assert "usage" not in tracer.trace_obj.spans[0].output
+
+
+class TestChatCompletionStream:
+    async def test_stream_requests_usage_and_span_output_carries_it(self):
+        captured_kwargs: dict = {}
+        chunks = [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk()]
+        service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs))
+
+        results = []
+        async for chunk in service.chat_completion_stream(
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True),
+            trace_id="trace-1",
+        ):
+            results.append(chunk)
+
+        assert len(results) == 3
+        assert captured_kwargs["stream_options"] == {"include_usage": True}
+        span = tracer.trace_obj.spans[0]
+        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert span.output["choices"][0]["message"]["content"] == "Hello!"
+
+
+class TestChatCompletionStreamAutoSend:
+    async def test_usage_only_final_chunk_reaches_span_output(self):
+        captured_kwargs: dict = {}
+        chunks = [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk()]
+        service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs))
+
+        await service.chat_completion_stream_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True),
+            trace_id="trace-1",
+        )
+
+        assert captured_kwargs["stream_options"] == {"include_usage": True}
+        span = tracer.trace_obj.spans[0]
+        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        # TaskMessage dump is still the base of the span output
+        assert span.output["id"] == "msg-1"
+
+    async def test_stream_without_usage_chunk_omits_usage(self):
+        captured_kwargs: dict = {}
+        chunks = [_delta_chunk("Hi", role="assistant")]
+        service, tracer = _make_service(_stream_gateway(chunks, captured_kwargs))
+
+        await service.chat_completion_stream_auto_send(
+            task_id="task-1",
+            llm_config=LLMConfig(model="gpt-4o", messages=[], stream=True),
+            trace_id="trace-1",
+        )
+
+        assert "usage" not in tracer.trace_obj.spans[0].output
diff --git a/tests/lib/utils/__init__.py b/tests/lib/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/lib/utils/test_completions.py b/tests/lib/utils/test_completions.py
new file mode 100644
index 000000000..3aa5f9120
--- /dev/null
+++ b/tests/lib/utils/test_completions.py
@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from agentex.lib.utils.completions import concat_completion_chunks
+from agentex.lib.types.llm_messages import Delta, Usage, Choice, Completion
+
+
+def _delta_chunk(content: str, role: str | None = None) -> Completion:
+    return Completion(choices=[Choice(index=0, delta=Delta(content=content, role=role))])
+
+
+def _usage_only_chunk(prompt: int, completion: int) -> Completion:
+    # stream_options.include_usage: litellm/OpenAI send a final chunk that has
+    # usage but an empty choices list
+    return Completion(
+        choices=[],
+        usage=Usage(prompt_tokens=prompt, completion_tokens=completion, total_tokens=prompt + completion),
+    )
+
+
+class TestConcatCompletionChunks:
+    def test_concatenates_delta_content(self):
+        result = concat_completion_chunks([_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!")])
+
+        assert result.choices[0].message.content == "Hello!"
+
+    def test_trailing_usage_only_chunk_keeps_choices_and_usage(self):
+        result = concat_completion_chunks(
+            [_delta_chunk("Hel", role="assistant"), _delta_chunk("lo!"), _usage_only_chunk(10, 5)]
+        )
+
+        assert result.choices[0].message.content == "Hello!"
+        assert result.usage is not None
+        assert result.usage.prompt_tokens == 10
+        assert result.usage.completion_tokens == 5
+        assert result.usage.total_tokens == 15
+
+    def test_usage_summed_across_chunks(self):
+        chunk_a = _delta_chunk("a", role="assistant")
+        chunk_a.usage = Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)
+        chunk_b = _delta_chunk("b")
+        chunk_b.usage = Usage(prompt_tokens=4, completion_tokens=5, total_tokens=9)
+
+        result = concat_completion_chunks([chunk_a, chunk_b])
+
+        assert result.usage.prompt_tokens == 5
+        assert result.usage.completion_tokens == 7
+        assert result.usage.total_tokens == 12
+
+    def test_no_usage_chunks_leave_usage_none(self):
+        result = concat_completion_chunks([_delta_chunk("x", role="assistant")])
+
+        assert result.usage is None

From ac466092117c9260671a28b77d6c8677beeb76ff Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 13:40:44 -0700
Subject: [PATCH 4/9] 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.
---
 .../10_async/00_base/030_tracing/README.md    |  31 +++++
 src/agentex/lib/adk/__init__.py               |   3 +-
 src/agentex/lib/adk/_modules/tracing.py       | 111 ++++++++++++++++++
 tests/lib/adk/test_tracing_module.py          |  94 ++++++++++++++-
 4 files changed, 237 insertions(+), 2 deletions(-)

diff --git a/examples/tutorials/10_async/00_base/030_tracing/README.md b/examples/tutorials/10_async/00_base/030_tracing/README.md
index 1d91f565f..c0601055e 100644
--- a/examples/tutorials/10_async/00_base/030_tracing/README.md
+++ b/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
diff --git a/src/agentex/lib/adk/__init__.py b/src/agentex/lib/adk/__init__.py
index 4d79be9dd..25b858485 100644
--- a/src/agentex/lib/adk/__init__.py
+++ b/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",
diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py
index 94bf741e4..9713c9e11 100644
--- a/src/agentex/lib/adk/_modules/tracing.py
+++ b/src/agentex/lib/adk/_modules/tracing.py
@@ -20,6 +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.types.span import Span
 from agentex.lib.utils.logging import make_logger
 from agentex.lib.utils.model_utils import BaseModel
@@ -42,6 +43,73 @@ 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: 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.
@@ -156,6 +224,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,
diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py
index 58d5d4a85..f13f4ccf0 100644
--- a/tests/lib/adk/test_tracing_module.py
+++ b/tests/lib/adk/test_tracing_module.py
@@ -8,7 +8,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
 
 
@@ -258,3 +258,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()

From 000ba5d7a0681a571049bf04d1024e3137a55a3c Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 13:49:37 -0700
Subject: [PATCH 5/9] test: fix sys.modules stub leakage and apply repo
 formatting

test_claude_agents_* now remove their placeholder packages after loading, which previously blocked real imports of the temporal plugin tree in later-collected tests. Formats touched files with ruff and narrows span.output types in new tests for pyright.
---
 src/agentex/lib/adk/_modules/tracing.py       |  4 +-
 .../models/temporal_tracing_model.py          | 69 +++++++++-----
 src/agentex/lib/utils/completions.py          | 14 +--
 tests/lib/adk/providers/test_litellm_usage.py | 22 +++--
 tests/lib/adk/test_tracing_module.py          |  4 +-
 .../plugins/openai_agents/test_model_usage.py | 90 ++++++++-----------
 tests/lib/test_claude_agents_activities.py    |  9 ++
 tests/lib/test_claude_agents_hooks.py         |  8 ++
 8 files changed, 120 insertions(+), 100 deletions(-)

diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py
index 9713c9e11..14219f28c 100644
--- a/src/agentex/lib/adk/_modules/tracing.py
+++ b/src/agentex/lib/adk/_modules/tracing.py
@@ -247,9 +247,7 @@ async def turn_span(
 
         Example::
 
-            async with adk.tracing.turn_span(
-                trace_id=task.id, name="turn", input={...}, task_id=task.id
-            ) as 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)
diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
index 02372aedd..11bb36cc7 100644
--- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
+++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
@@ -7,6 +7,7 @@
 The key innovation is that these are thin wrappers around the standard OpenAI models,
 avoiding code duplication while adding tracing capabilities.
 """
+
 from __future__ import annotations
 
 import logging
@@ -51,24 +52,28 @@ def _serialize_item(item: Any) -> dict[str, Any]:
     Uses model_dump() for Pydantic models, otherwise extracts attributes manually.
     Filters out internal Pydantic fields that can't be serialized.
     """
-    if hasattr(item, 'model_dump'):
+    if hasattr(item, "model_dump"):
         # Pydantic model - use model_dump for proper serialization
         try:
-            return item.model_dump(mode='json', exclude_unset=True)
+            return item.model_dump(mode="json", exclude_unset=True)
         except Exception:
             # Fallback to dict conversion
-            return dict(item) if hasattr(item, '__iter__') else {}
+            return dict(item) if hasattr(item, "__iter__") else {}
     else:
         # Not a Pydantic model - extract attributes manually
         item_dict = {}
         for attr_name in dir(item):
-            if not attr_name.startswith('_') and attr_name not in ('model_fields', 'model_config', 'model_computed_fields'):
+            if not attr_name.startswith("_") and attr_name not in (
+                "model_fields",
+                "model_config",
+                "model_computed_fields",
+            ):
                 try:
                     attr_value = getattr(item, attr_name, None)
                     # Skip methods and None values
                     if attr_value is not None and not callable(attr_value):
                         # Convert to JSON-serializable format
-                        if hasattr(attr_value, 'model_dump'):
+                        if hasattr(attr_value, "model_dump"):
                             item_dict[attr_name] = attr_value.model_dump()
                         elif isinstance(attr_value, (str, int, float, bool, list, dict)):
                             item_dict[attr_name] = attr_value
@@ -106,7 +111,9 @@ def __init__(self, openai_client: Optional[AsyncOpenAI] = None, **kwargs):
         # Initialize tracer for all models
         agentex_client = create_async_agentex_client()
         self._tracer = AsyncTracer(agentex_client)
-        logger.info(f"[TemporalTracingModelProvider] Initialized with AgentEx tracer, custom_client={openai_client is not None}")
+        logger.info(
+            f"[TemporalTracingModelProvider] Initialized with AgentEx tracer, custom_client={openai_client is not None}"
+        )
 
     @override
     def get_model(self, model_name: Optional[str]) -> Model:
@@ -126,10 +133,14 @@ def get_model(self, model_name: Optional[str]) -> Model:
             logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIResponsesModel '{model_name}' with tracing")
             return TemporalTracingResponsesModel(base_model, self._tracer)  # type: ignore[abstract]
         elif isinstance(base_model, OpenAIChatCompletionsModel):
-            logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIChatCompletionsModel '{model_name}' with tracing")
+            logger.info(
+                f"[TemporalTracingModelProvider] Wrapping OpenAIChatCompletionsModel '{model_name}' with tracing"
+            )
             return TemporalTracingChatCompletionsModel(base_model, self._tracer)  # type: ignore[abstract]
         else:
-            logger.warning(f"[TemporalTracingModelProvider] Unknown model type, returning without tracing: {type(base_model)}")
+            logger.warning(
+                f"[TemporalTracingModelProvider] Unknown model type, returning without tracing: {type(base_model)}"
+            )
             return base_model
 
 
@@ -181,7 +192,9 @@ async def get_response(
 
         # If we have tracing context, wrap with span
         if trace_id and parent_span_id:
-            logger.debug(f"[TemporalTracingResponsesModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}")
+            logger.debug(
+                f"[TemporalTracingResponsesModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}"
+            )
 
             trace = self._tracer.trace(trace_id)
 
@@ -199,7 +212,9 @@ async def get_response(
                         "temperature": model_settings.temperature,
                         "max_tokens": model_settings.max_tokens,
                         "reasoning": model_settings.reasoning,
-                    } if model_settings else None,
+                    }
+                    if model_settings
+                    else None,
                 },
             ) as span:
                 try:
@@ -222,7 +237,7 @@ async def get_response(
                     new_items = []
                     final_output = None
 
-                    if hasattr(response, 'output') and response.output:
+                    if hasattr(response, "output") and response.output:
                         response_output = response.output if isinstance(response.output, list) else [response.output]
 
                         for item in response_output:
@@ -232,12 +247,12 @@ async def get_response(
                                     new_items.append(item_dict)
 
                                     # Extract final_output from message type if available
-                                    if item_dict.get('type') == 'message' and not final_output:
-                                        content = item_dict.get('content', [])
+                                    if item_dict.get("type") == "message" and not final_output:
+                                        content = item_dict.get("content", [])
                                         if content and isinstance(content, list):
                                             for content_part in content:
-                                                if isinstance(content_part, dict) and 'text' in content_part:
-                                                    final_output = content_part['text']
+                                                if isinstance(content_part, dict) and "text" in content_part:
+                                                    final_output = content_part["text"]
                                                     break
                             except Exception as e:
                                 logger.warning(f"Failed to serialize item in temporal tracing model: {e}")
@@ -329,7 +344,9 @@ async def get_response(
 
         # If we have tracing context, wrap with span
         if trace_id and parent_span_id:
-            logger.debug(f"[TemporalTracingChatCompletionsModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}")
+            logger.debug(
+                f"[TemporalTracingChatCompletionsModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}"
+            )
 
             trace = self._tracer.trace(trace_id)
 
@@ -346,7 +363,9 @@ async def get_response(
                     "model_settings": {
                         "temperature": model_settings.temperature,
                         "max_tokens": model_settings.max_tokens,
-                    } if model_settings else None,
+                    }
+                    if model_settings
+                    else None,
                 },
             ) as span:
                 try:
@@ -366,7 +385,7 @@ async def get_response(
                     new_items = []
                     final_output = None
 
-                    if hasattr(response, 'output') and response.output:
+                    if hasattr(response, "output") and response.output:
                         response_output = response.output if isinstance(response.output, list) else [response.output]
 
                         for item in response_output:
@@ -376,12 +395,12 @@ async def get_response(
                                     new_items.append(item_dict)
 
                                     # Extract final_output from message type if available
-                                    if item_dict.get('type') == 'message' and not final_output:
-                                        content = item_dict.get('content', [])
+                                    if item_dict.get("type") == "message" and not final_output:
+                                        content = item_dict.get("content", [])
                                         if content and isinstance(content, list):
                                             for content_part in content:
-                                                if isinstance(content_part, dict) and 'text' in content_part:
-                                                    final_output = content_part['text']
+                                                if isinstance(content_part, dict) and "text" in content_part:
+                                                    final_output = content_part["text"]
                                                     break
                             except Exception as e:
                                 logger.warning(f"Failed to serialize item in temporal tracing model: {e}")
@@ -406,7 +425,9 @@ async def get_response(
                     raise
         else:
             # No tracing context, just pass through
-            logger.debug("[TemporalTracingChatCompletionsModel] No tracing context available, calling base model directly")
+            logger.debug(
+                "[TemporalTracingChatCompletionsModel] No tracing context available, calling base model directly"
+            )
             return await self._base_model.get_response(
                 system_instructions=system_instructions,
                 input=input,
@@ -422,4 +443,4 @@ async def get_response(
     def stream_response(self, *args, **kwargs):
         """Streaming is handled via get_response in Temporal activities.
         Required so the class is concrete and instantiable at runtime."""
-        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
\ No newline at end of file
+        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
diff --git a/src/agentex/lib/utils/completions.py b/src/agentex/lib/utils/completions.py
index 595e0fb01..3cb1d7b4d 100644
--- a/src/agentex/lib/utils/completions.py
+++ b/src/agentex/lib/utils/completions.py
@@ -26,8 +26,7 @@ def _(a: Completion, b: Completion) -> Completion:
     # final chunk carries usage but no choices. Keep the unpaired side instead
     # of truncating to the shorter list.
     a.choices = [
-        x if y is None else y if x is None else _concat_chunks(x, y)
-        for x, y in zip_longest(a.choices, b.choices)
+        x if y is None else y if x is None else _concat_chunks(x, y) for x, y in zip_longest(a.choices, b.choices)
     ]
     a.usage = _concat_chunks(a.usage, b.usage)
 
@@ -45,6 +44,7 @@ def _(a: Choice, b: Choice) -> Choice:
     a.finish_reason = a.finish_reason or b.finish_reason
     return a
 
+
 @_concat_chunks.register
 def _(a: Usage | None, b: Usage | None) -> Usage | None:
     if a is not None and b is not None:
@@ -68,9 +68,7 @@ def _(a: Delta, b: Delta) -> Delta:
             if tool_call.index not in grouped_tool_calls:
                 grouped_tool_calls[tool_call.index] = tool_call
             else:
-                grouped_tool_calls[tool_call.index] = _concat_chunks(
-                    grouped_tool_calls[tool_call.index], tool_call
-                )
+                grouped_tool_calls[tool_call.index] = _concat_chunks(grouped_tool_calls[tool_call.index], tool_call)
 
         a.tool_calls = list(grouped_tool_calls.values())
     elif hasattr(b, "tool_calls") and b.tool_calls:
@@ -88,11 +86,7 @@ def _(a: ToolCallRequest, b: ToolCallRequest) -> ToolCallRequest:
     index_val = a.index if hasattr(a, "index") and a.index is not None else b.index
 
     # Concatenate the function part
-    function_val = (
-        _concat_chunks(a.function, b.function)
-        if a.function and b.function
-        else a.function or b.function
-    )
+    function_val = _concat_chunks(a.function, b.function) if a.function and b.function else a.function or b.function
 
     # Set all properties
     a.id = id_val
diff --git a/tests/lib/adk/providers/test_litellm_usage.py b/tests/lib/adk/providers/test_litellm_usage.py
index 5f3bee405..5f5d480d9 100644
--- a/tests/lib/adk/providers/test_litellm_usage.py
+++ b/tests/lib/adk/providers/test_litellm_usage.py
@@ -7,13 +7,13 @@
 
 from __future__ import annotations
 
+from typing import Any
 from datetime import UTC, datetime
 from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, MagicMock
 
 from agentex.types.span import Span
 from agentex.types.task_message import TaskMessage
-from agentex.types.task_message_content import TextContent
 from agentex.lib.types.llm_messages import (
     Delta,
     Usage,
@@ -22,6 +22,7 @@
     Completion,
     AssistantMessage,
 )
+from agentex.types.task_message_content import TextContent
 from agentex.lib.core.services.adk.providers.litellm import (
     LiteLLMService,
     _stream_kwargs_with_usage,
@@ -54,6 +55,11 @@ def trace(self, trace_id):
         return self.trace_obj
 
 
+def _output_dict(span: Span) -> dict[str, Any]:
+    assert isinstance(span.output, dict)
+    return span.output
+
+
 def _make_streaming_service():
     streaming_context = MagicMock()
     streaming_context.task_message = TaskMessage(
@@ -142,7 +148,7 @@ async def test_span_output_carries_usage(self):
         )
 
         span = tracer.trace_obj.spans[0]
-        assert span.output["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
+        assert _output_dict(span)["usage"] == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
 
     async def test_span_output_omits_usage_when_absent(self):
         completion = Completion(
@@ -158,7 +164,7 @@ async def test_span_output_omits_usage_when_absent(self):
             trace_id="trace-1",
         )
 
-        assert "usage" not in tracer.trace_obj.spans[0].output
+        assert "usage" not in _output_dict(tracer.trace_obj.spans[0])
 
 
 class TestChatCompletionStream:
@@ -177,8 +183,8 @@ async def test_stream_requests_usage_and_span_output_carries_it(self):
         assert len(results) == 3
         assert captured_kwargs["stream_options"] == {"include_usage": True}
         span = tracer.trace_obj.spans[0]
-        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
-        assert span.output["choices"][0]["message"]["content"] == "Hello!"
+        assert _output_dict(span)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert _output_dict(span)["choices"][0]["message"]["content"] == "Hello!"
 
 
 class TestChatCompletionStreamAutoSend:
@@ -195,9 +201,9 @@ async def test_usage_only_final_chunk_reaches_span_output(self):
 
         assert captured_kwargs["stream_options"] == {"include_usage": True}
         span = tracer.trace_obj.spans[0]
-        assert span.output["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
+        assert _output_dict(span)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
         # TaskMessage dump is still the base of the span output
-        assert span.output["id"] == "msg-1"
+        assert _output_dict(span)["id"] == "msg-1"
 
     async def test_stream_without_usage_chunk_omits_usage(self):
         captured_kwargs: dict = {}
@@ -210,4 +216,4 @@ async def test_stream_without_usage_chunk_omits_usage(self):
             trace_id="trace-1",
         )
 
-        assert "usage" not in tracer.trace_obj.spans[0].output
+        assert "usage" not in _output_dict(tracer.trace_obj.spans[0])
diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py
index f13f4ccf0..773fc00b4 100644
--- a/tests/lib/adk/test_tracing_module.py
+++ b/tests/lib/adk/test_tracing_module.py
@@ -314,9 +314,7 @@ async def test_turn_span_preserves_existing_data(self):
         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:
+            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"]
diff --git a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
index 26db6fb5f..4bd6a0ee1 100644
--- a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
+++ b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
@@ -8,15 +8,32 @@
 
 from __future__ import annotations
 
+from typing import Any
 from datetime import UTC, datetime
 from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from agents import ModelResponse
-from agents.usage import Usage, InputTokensDetails, OutputTokensDetails
-
+from agents import ModelResponse, ModelSettings
+from agents.usage import Usage
+from openai.types.responses import Response, ResponseCompletedEvent
+from openai.types.responses.response_usage import (
+    ResponseUsage,
+    InputTokensDetails,
+    OutputTokensDetails,
+)
+
+import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm
 from agentex.types.span import Span
+from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
+    TemporalTracingResponsesModel,
+    TemporalTracingChatCompletionsModel,
+)
+from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
+    streaming_task_id,
+    streaming_trace_id,
+    streaming_parent_span_id,
+)
 
 pytestmark = pytest.mark.asyncio
 
@@ -53,12 +70,6 @@ def trace(self, trace_id):
 
 @pytest.fixture
 def tracing_contextvars():
-    from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
-        streaming_task_id,
-        streaming_trace_id,
-        streaming_parent_span_id,
-    )
-
     tokens = [
         streaming_task_id.set("task-1"),
         streaming_trace_id.set("trace-1"),
@@ -80,6 +91,7 @@ def _agents_usage() -> Usage:
         output_tokens_details=OutputTokensDetails(reasoning_tokens=40),
     )
 
+
 EXPECTED_USAGE_BLOB = {
     "input_tokens": 120,
     "output_tokens": 80,
@@ -89,6 +101,11 @@ def _agents_usage() -> Usage:
 }
 
 
+def _output_dict(span: Span) -> dict[str, Any]:
+    assert isinstance(span.output, dict)
+    return span.output
+
+
 class TestTemporalTracingModels:
     async def _run_wrapper(self, wrapper_cls) -> Span:
         tracer = FakeTracer()
@@ -99,8 +116,6 @@ async def _run_wrapper(self, wrapper_cls) -> Span:
         )
 
         model = wrapper_cls(base_model, tracer)
-        from agents import ModelSettings
-
         response = await model.get_response(
             system_instructions=None,
             input="hello",
@@ -115,20 +130,12 @@ async def _run_wrapper(self, wrapper_cls) -> Span:
         return tracer.trace_obj.spans[0]
 
     async def test_responses_model_writes_usage_to_span_output(self, tracing_contextvars):
-        from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
-            TemporalTracingResponsesModel,
-        )
-
         span = await self._run_wrapper(TemporalTracingResponsesModel)
-        assert span.output["usage"] == EXPECTED_USAGE_BLOB
+        assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
 
     async def test_chat_completions_model_writes_usage_to_span_output(self, tracing_contextvars):
-        from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
-            TemporalTracingChatCompletionsModel,
-        )
-
         span = await self._run_wrapper(TemporalTracingChatCompletionsModel)
-        assert span.output["usage"] == EXPECTED_USAGE_BLOB
+        assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
 
 
 class FakeStream:
@@ -145,21 +152,12 @@ async def gen():
 
 class TestTemporalStreamingModel:
     async def test_streaming_model_captures_final_response_usage(self, tracing_contextvars):
-        import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm
-
-        from openai.types.responses import Response, ResponseCompletedEvent
-        from openai.types.responses.response_usage import (
-            ResponseUsage,
-            InputTokensDetails as ResponseInputTokensDetails,
-            OutputTokensDetails as ResponseOutputTokensDetails,
-        )
-
         usage = ResponseUsage(
             input_tokens=120,
             output_tokens=80,
             total_tokens=200,
-            input_tokens_details=ResponseInputTokensDetails(cached_tokens=30),
-            output_tokens_details=ResponseOutputTokensDetails(reasoning_tokens=40),
+            input_tokens_details=InputTokensDetails(cached_tokens=30),
+            output_tokens_details=OutputTokensDetails(reasoning_tokens=40),
         )
         completed = ResponseCompletedEvent.model_construct(
             type="response.completed",
@@ -170,13 +168,9 @@ async def test_streaming_model_captures_final_response_usage(self, tracing_conte
         openai_client = MagicMock()
         openai_client.responses.create = AsyncMock(return_value=FakeStream([completed]))
 
-        with (
-            patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()),
-            patch.object(tsm, "AsyncTracer", return_value=fake_tracer),
-        ):
-            model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client)
-
-        from agents import ModelSettings
+        with patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()):
+            with patch.object(tsm, "AsyncTracer", return_value=fake_tracer):
+                model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client)
 
         response = await model.get_response(
             system_instructions=None,
@@ -199,13 +193,9 @@ async def test_streaming_model_captures_final_response_usage(self, tracing_conte
         # And on the span output for billing
         assert len(fake_tracer.trace_obj.spans) == 1
         span = fake_tracer.trace_obj.spans[0]
-        assert span.output["usage"] == EXPECTED_USAGE_BLOB
+        assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
 
     async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_contextvars):
-        import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm
-
-        from openai.types.responses import Response, ResponseCompletedEvent
-
         completed = ResponseCompletedEvent.model_construct(
             type="response.completed",
             response=Response.model_construct(output=[], usage=None),
@@ -215,13 +205,9 @@ async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_c
         openai_client = MagicMock()
         openai_client.responses.create = AsyncMock(return_value=FakeStream([completed]))
 
-        with (
-            patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()),
-            patch.object(tsm, "AsyncTracer", return_value=fake_tracer),
-        ):
-            model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client)
-
-        from agents import ModelSettings
+        with patch.object(tsm, "create_async_agentex_client", return_value=MagicMock()):
+            with patch.object(tsm, "AsyncTracer", return_value=fake_tracer):
+                model = tsm.TemporalStreamingModel(model_name="gpt-4o", openai_client=openai_client)
 
         response = await model.get_response(
             system_instructions=None,
@@ -235,4 +221,4 @@ async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_c
 
         assert response.usage.input_tokens == 0
         span = fake_tracer.trace_obj.spans[0]
-        assert "usage" not in span.output
+        assert "usage" not in _output_dict(span)
diff --git a/tests/lib/test_claude_agents_activities.py b/tests/lib/test_claude_agents_activities.py
index fa633cd6b..99982d851 100644
--- a/tests/lib/test_claude_agents_activities.py
+++ b/tests/lib/test_claude_agents_activities.py
@@ -59,6 +59,7 @@
     sys.modules.setdefault(_name, _mock)
 
 # Also ensure parent packages exist as stubs so Python resolves the dotted path
+_created_pkg_stubs: list[str] = []
 for _pkg in [
     "agentex.lib.core.temporal.plugins",
     "agentex.lib.core.temporal.plugins.claude_agents",
@@ -70,6 +71,7 @@
         _mod.__path__ = []  # type: ignore[attr-defined]
         _mod.__package__ = _pkg
         sys.modules[_pkg] = _mod
+        _created_pkg_stubs.append(_pkg)
 
 # Load activities.py directly from its file path
 _spec = importlib.util.spec_from_file_location(
@@ -81,6 +83,13 @@
 sys.modules[_spec.name] = _activities_mod
 _spec.loader.exec_module(_activities_mod)
 
+# Drop the placeholder packages now that the module is loaded: their empty
+# __path__ would otherwise block real imports of agentex.lib.core.temporal.*
+# submodules in every test collected after this file. The loaded module keeps
+# its bindings.
+for _pkg in _created_pkg_stubs:
+    sys.modules.pop(_pkg, None)
+
 _reconstruct_agent_defs = _activities_mod._reconstruct_agent_defs  # type: ignore[attr-defined]
 claude_options_to_dict = _activities_mod.claude_options_to_dict  # type: ignore[attr-defined]
 
diff --git a/tests/lib/test_claude_agents_hooks.py b/tests/lib/test_claude_agents_hooks.py
index 373956680..ec82b16db 100644
--- a/tests/lib/test_claude_agents_hooks.py
+++ b/tests/lib/test_claude_agents_hooks.py
@@ -29,6 +29,7 @@
     sys.modules.setdefault(_name, MagicMock())
 
 # Ensure parent packages exist so the dotted path resolves
+_created_pkg_stubs: list[str] = []
 for _pkg in [
     "agentex",
     "agentex.lib",
@@ -43,6 +44,7 @@
         _mod.__path__ = []  # type: ignore[attr-defined]
         _mod.__package__ = _pkg
         sys.modules[_pkg] = _mod
+        _created_pkg_stubs.append(_pkg)
 
 # Real pydantic types — importable without triggering the problematic chain.
 from agentex.types.tool_response_content import ToolResponseContent  # noqa: E402
@@ -57,6 +59,12 @@
 sys.modules[_spec.name] = _hooks_mod
 _spec.loader.exec_module(_hooks_mod)
 
+# Drop the placeholder packages now that the module is loaded: their empty
+# __path__ would otherwise block real imports of agentex.* submodules in every
+# test collected after this file. The loaded module keeps its bindings.
+for _pkg in _created_pkg_stubs:
+    sys.modules.pop(_pkg, None)
+
 TemporalStreamingHooks = _hooks_mod.TemporalStreamingHooks
 create_streaming_hooks = _hooks_mod.create_streaming_hooks
 

From 06c9b9c64194407a302eb215c49896d4d6d1cd9a Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 14:37:37 -0700
Subject: [PATCH 6/9] fix(adk): warn when record_usage replaces non-dict span
 data

---
 src/agentex/lib/adk/_modules/tracing.py |  5 +++++
 tests/lib/adk/test_tracing_module.py    | 28 +++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py
index 14219f28c..be30068a0 100644
--- a/src/agentex/lib/adk/_modules/tracing.py
+++ b/src/agentex/lib/adk/_modules/tracing.py
@@ -93,6 +93,11 @@ def record_usage(
             )
         )
 
+        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
diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py
index 773fc00b4..4d1fc1f23 100644
--- a/tests/lib/adk/test_tracing_module.py
+++ b/tests/lib/adk/test_tracing_module.py
@@ -321,6 +321,34 @@ async def test_turn_span_preserves_existing_data(self):
         assert ended_span.data["custom"] == "value"
         assert ended_span.data["usage"] == {"prompt_tokens": 3, "completion_tokens": 4}
 
+    async def test_turn_span_warns_and_replaces_non_dict_data(self, caplog):
+        mock_service, module = _make_module()
+        started = _make_span(data=[{"item": 1}])
+        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={"input_tokens": 1, "output_tokens": 2})
+
+        assert any("existing data will be replaced" in message for message in caplog.messages)
+        ended_span = mock_service.end_span.call_args.kwargs["span"]
+        assert ended_span.data == {"usage": {"input_tokens": 1, "output_tokens": 2}}
+
+    async def test_turn_span_dict_data_does_not_warn(self, caplog):
+        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") as turn:
+                with caplog.at_level("WARNING"):
+                    turn.record_usage(usage={"input_tokens": 1})
+
+        assert not any("existing data will be replaced" in message for message in caplog.messages)
+
     async def test_turn_span_cost_only(self):
         mock_service, module = _make_module()
         started = _make_span()

From 96b28592f0b66cb887ca3143e6cd602a474ab716 Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 15:05:20 -0700
Subject: [PATCH 7/9] test: adapt streaming model usage tests to next's
 always-emit behavior

---
 tests/lib/adk/test_tracing_module.py          | 57 +++++++++----------
 .../plugins/openai_agents/test_model_usage.py | 13 ++++-
 2 files changed, 38 insertions(+), 32 deletions(-)

diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py
index 4d1fc1f23..adb8641d5 100644
--- a/tests/lib/adk/test_tracing_module.py
+++ b/tests/lib/adk/test_tracing_module.py
@@ -113,10 +113,11 @@ async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
         mock_service, module = _make_module()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger") as mock_logger, patch.object(
+            _tracing_mod.workflow, "metric_meter", return_value=mock_meter
+        ):
             mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
             result = await module.start_span(trace_id="trace-123", name="test-span")
 
@@ -127,9 +128,7 @@ async def test_start_span_in_workflow_returns_none_when_activity_fails(self):
             description="Temporal tracing span activities dropped after fail-open",
             unit="1",
         )
-        mock_meter.create_counter.return_value.add.assert_called_once_with(
-            1, {"event_type": "start"}
-        )
+        mock_meter.create_counter.return_value.add.assert_called_once_with(1, {"event_type": "start"})
         mock_helpers.execute_activity.assert_called_once()
         mock_service.start_span.assert_not_called()
 
@@ -138,10 +137,11 @@ async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
         span = _make_span()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger") as mock_logger, patch.object(
+            _tracing_mod.workflow, "metric_meter", return_value=mock_meter
+        ):
             mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
             result = await module.end_span(trace_id="trace-123", span=span)
 
@@ -152,18 +152,16 @@ async def test_end_span_in_workflow_returns_span_when_activity_fails(self):
             description="Temporal tracing span activities dropped after fail-open",
             unit="1",
         )
-        mock_meter.create_counter.return_value.add.assert_called_once_with(
-            1, {"event_type": "end"}
-        )
+        mock_meter.create_counter.return_value.add.assert_called_once_with(1, {"event_type": "end"})
         mock_helpers.execute_activity.assert_called_once()
         mock_service.end_span.assert_not_called()
 
     async def test_context_manager_skips_end_when_temporal_start_fails(self):
         mock_service, module = _make_module()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod.workflow, "logger"):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod.workflow, "logger"):
             mock_helpers.execute_activity = AsyncMock(side_effect=_make_activity_error())
             async with module.span(trace_id="trace-123", name="test-span") as span:
                 assert span is None
@@ -175,8 +173,9 @@ async def test_context_manager_skips_end_when_temporal_start_fails(self):
     async def test_start_span_in_workflow_propagates_unexpected_errors(self):
         mock_service, module = _make_module()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers:
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers:
             mock_helpers.execute_activity = AsyncMock(side_effect=RuntimeError("bad response shape"))
             try:
                 await module.start_span(trace_id="trace-123", name="test-span")
@@ -193,11 +192,11 @@ async def test_start_span_in_workflow_propagates_cancellation(self):
         activity_error = _make_activity_error()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), patch.object(
+            _tracing_mod.workflow, "logger"
+        ) as mock_logger, patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
             mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
 
             with pytest.raises(ActivityError):
@@ -214,11 +213,11 @@ async def test_end_span_in_workflow_propagates_cancellation(self):
         activity_error = _make_activity_error()
         mock_meter = _make_metric_meter()
 
-        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), \
-                patch.object(_tracing_mod, "ActivityHelpers") as mock_helpers, \
-                patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), \
-                patch.object(_tracing_mod.workflow, "logger") as mock_logger, \
-                patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
+        with patch.object(_tracing_mod, "in_temporal_workflow", return_value=True), patch.object(
+            _tracing_mod, "ActivityHelpers"
+        ) as mock_helpers, patch.object(_tracing_mod, "is_cancelled_exception", return_value=True), patch.object(
+            _tracing_mod.workflow, "logger"
+        ) as mock_logger, patch.object(_tracing_mod.workflow, "metric_meter", return_value=mock_meter):
             mock_helpers.execute_activity = AsyncMock(side_effect=activity_error)
 
             with pytest.raises(ActivityError):
diff --git a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
index 4bd6a0ee1..282b95e26 100644
--- a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
+++ b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
@@ -183,7 +183,6 @@ async def test_streaming_model_captures_final_response_usage(self, tracing_conte
         )
 
         # Real usage lands on the returned ModelResponse (was zeroed before)
-        assert response.usage.requests == 1
         assert response.usage.input_tokens == 120
         assert response.usage.output_tokens == 80
         assert response.usage.total_tokens == 200
@@ -195,7 +194,7 @@ async def test_streaming_model_captures_final_response_usage(self, tracing_conte
         span = fake_tracer.trace_obj.spans[0]
         assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
 
-    async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_contextvars):
+    async def test_streaming_model_writes_zero_usage_when_api_reports_none(self, tracing_contextvars):
         completed = ResponseCompletedEvent.model_construct(
             type="response.completed",
             response=Response.model_construct(output=[], usage=None),
@@ -219,6 +218,14 @@ async def test_streaming_model_omits_usage_when_api_reports_none(self, tracing_c
             tracing=None,
         )
 
+        # No usage from the API: the model reports zeros rather than omitting,
+        # so billing sums 0 instead of missing the span
         assert response.usage.input_tokens == 0
         span = fake_tracer.trace_obj.spans[0]
-        assert "usage" not in _output_dict(span)
+        assert _output_dict(span)["usage"] == {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "total_tokens": 0,
+            "cached_input_tokens": 0,
+            "reasoning_tokens": 0,
+        }

From 4b99087a2a7b8a4ab4f1b2013353d12adb22c5d6 Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 15:14:52 -0700
Subject: [PATCH 8/9] refactor(tracing): remove dead TemporalTracing wrapper
 models

The wrappers never implemented abstract stream_response, so TemporalTracingModelProvider.get_model() raised TypeError on every call since introduction; no working callers can exist. The streaming provider plus run.py hooks are the live tracing path.
---
 .../lib/core/temporal/plugins/__init__.py     |   3 -
 .../plugins/openai_agents/__init__.py         |   4 -
 .../plugins/openai_agents/models/__init__.py  |   8 -
 .../models/temporal_tracing_model.py          | 446 ------------------
 .../plugins/openai_agents/test_model_usage.py |  57 +--
 5 files changed, 4 insertions(+), 514 deletions(-)
 delete mode 100644 src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py

diff --git a/src/agentex/lib/core/temporal/plugins/__init__.py b/src/agentex/lib/core/temporal/plugins/__init__.py
index 52ab6eac7..8377ab5fc 100644
--- a/src/agentex/lib/core/temporal/plugins/__init__.py
+++ b/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,7 +46,6 @@
 __all__ = [
     "TemporalStreamingModel",
     "TemporalStreamingModelProvider",
-    "TemporalTracingModelProvider",
     "ContextInterceptor",
     "streaming_task_id",
     "streaming_trace_id",
diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py b/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py
index 7d81b37d0..453e27074 100644
--- a/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py
+++ b/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",
diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py
index bb5dc97ed..6f972a2ed 100644
--- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py
+++ b/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",
 ]
\ No newline at end of file
diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
deleted file mode 100644
index 11bb36cc7..000000000
--- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_tracing_model.py
+++ /dev/null
@@ -1,446 +0,0 @@
-"""Temporal-aware tracing model provider.
-
-This module provides model implementations that add AgentEx tracing to standard OpenAI models
-when running in Temporal workflows/activities. It uses context variables set by the Temporal
-context interceptor to access task_id, trace_id, and parent_span_id.
-
-The key innovation is that these are thin wrappers around the standard OpenAI models,
-avoiding code duplication while adding tracing capabilities.
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import Any, List, Union, Optional, override
-
-from agents import (
-    Tool,
-    Model,
-    Handoff,
-    ModelTracing,
-    ModelResponse,
-    ModelSettings,
-    OpenAIProvider,
-    TResponseInputItem,
-    AgentOutputSchemaBase,
-)
-from openai import AsyncOpenAI
-from openai.types.responses import ResponsePromptParam
-from agents.models.openai_responses import OpenAIResponsesModel
-from agents.models.openai_chatcompletions import OpenAIChatCompletionsModel
-
-from agentex.lib.core.tracing.usage import usage_from_openai_response_usage
-from agentex.lib.core.tracing.tracer import AsyncTracer
-
-# Import AgentEx components
-from agentex.lib.adk.utils._modules.client import create_async_agentex_client
-
-# Import context variables from the interceptor
-from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
-    streaming_task_id,
-    streaming_trace_id,
-    streaming_parent_span_id,
-)
-
-logger = logging.getLogger("agentex.temporal.tracing")
-
-
-def _serialize_item(item: Any) -> dict[str, Any]:
-    """
-    Universal serializer for any item type from OpenAI Agents SDK.
-
-    Uses model_dump() for Pydantic models, otherwise extracts attributes manually.
-    Filters out internal Pydantic fields that can't be serialized.
-    """
-    if hasattr(item, "model_dump"):
-        # Pydantic model - use model_dump for proper serialization
-        try:
-            return item.model_dump(mode="json", exclude_unset=True)
-        except Exception:
-            # Fallback to dict conversion
-            return dict(item) if hasattr(item, "__iter__") else {}
-    else:
-        # Not a Pydantic model - extract attributes manually
-        item_dict = {}
-        for attr_name in dir(item):
-            if not attr_name.startswith("_") and attr_name not in (
-                "model_fields",
-                "model_config",
-                "model_computed_fields",
-            ):
-                try:
-                    attr_value = getattr(item, attr_name, None)
-                    # Skip methods and None values
-                    if attr_value is not None and not callable(attr_value):
-                        # Convert to JSON-serializable format
-                        if hasattr(attr_value, "model_dump"):
-                            item_dict[attr_name] = attr_value.model_dump()
-                        elif isinstance(attr_value, (str, int, float, bool, list, dict)):
-                            item_dict[attr_name] = attr_value
-                        else:
-                            item_dict[attr_name] = str(attr_value)
-                except Exception:
-                    # Skip attributes that can't be accessed
-                    pass
-        return item_dict
-
-
-class TemporalTracingModelProvider(OpenAIProvider):
-    """Model provider that returns OpenAI models wrapped with AgentEx tracing.
-
-    This provider extends the standard OpenAIProvider to return models that add
-    tracing spans around model calls when running in Temporal activities with
-    the context interceptor enabled.
-    """
-
-    def __init__(self, openai_client: Optional[AsyncOpenAI] = None, **kwargs):
-        """Initialize the tracing model provider.
-
-        Args:
-            openai_client: Optional custom AsyncOpenAI client. If provided, this client
-                          will be used for all model calls. If not provided, OpenAIProvider
-                          will create a default client.
-            **kwargs: All other arguments are passed to OpenAIProvider.
-        """
-        # Pass openai_client to parent if provided
-        if openai_client is not None:
-            super().__init__(openai_client=openai_client, **kwargs)
-        else:
-            super().__init__(**kwargs)
-
-        # Initialize tracer for all models
-        agentex_client = create_async_agentex_client()
-        self._tracer = AsyncTracer(agentex_client)
-        logger.info(
-            f"[TemporalTracingModelProvider] Initialized with AgentEx tracer, custom_client={openai_client is not None}"
-        )
-
-    @override
-    def get_model(self, model_name: Optional[str]) -> Model:
-        """Get a model wrapped with tracing capabilities.
-
-        Args:
-            model_name: The name of the model to use
-
-        Returns:
-            A model instance wrapped with tracing
-        """
-        # Get the base model from the parent provider
-        base_model = super().get_model(model_name)
-
-        # Wrap with appropriate tracing wrapper based on model type
-        if isinstance(base_model, OpenAIResponsesModel):
-            logger.info(f"[TemporalTracingModelProvider] Wrapping OpenAIResponsesModel '{model_name}' with tracing")
-            return TemporalTracingResponsesModel(base_model, self._tracer)  # type: ignore[abstract]
-        elif isinstance(base_model, OpenAIChatCompletionsModel):
-            logger.info(
-                f"[TemporalTracingModelProvider] Wrapping OpenAIChatCompletionsModel '{model_name}' with tracing"
-            )
-            return TemporalTracingChatCompletionsModel(base_model, self._tracer)  # type: ignore[abstract]
-        else:
-            logger.warning(
-                f"[TemporalTracingModelProvider] Unknown model type, returning without tracing: {type(base_model)}"
-            )
-            return base_model
-
-
-class TemporalTracingResponsesModel(Model):
-    """Wrapper for OpenAIResponsesModel that adds AgentEx tracing.
-
-    This is a thin wrapper that adds tracing spans around the base model's
-    get_response() method. It reads tracing context from ContextVars set by
-    the Temporal context interceptor.
-    """
-
-    def __init__(self, base_model: OpenAIResponsesModel, tracer: AsyncTracer):
-        """Initialize the tracing wrapper.
-
-        Args:
-            base_model: The OpenAI Responses model to wrap
-            tracer: The AgentEx tracer to use
-        """
-        self._base_model = base_model
-        self._tracer = tracer
-        # Expose the model name for compatibility
-        self.model = base_model.model
-
-    @override
-    async def get_response(
-        self,
-        system_instructions: Optional[str],
-        input: Union[str, List[TResponseInputItem]],
-        model_settings: ModelSettings,
-        tools: List[Tool],
-        output_schema: Optional[AgentOutputSchemaBase],
-        handoffs: List[Handoff],
-        tracing: ModelTracing,
-        previous_response_id: Optional[str] = None,
-        conversation_id: Optional[str] = None,
-        prompt: Optional[ResponsePromptParam] = None,
-        **kwargs,
-    ) -> ModelResponse:
-        """Get a response from the model with optional tracing.
-
-        If tracing context is available from the interceptor, this wraps the
-        model call in a tracing span. Otherwise, it passes through to the
-        base model without tracing.
-        """
-        # Try to get tracing context from ContextVars
-        task_id = streaming_task_id.get()
-        trace_id = streaming_trace_id.get()
-        parent_span_id = streaming_parent_span_id.get()
-
-        # If we have tracing context, wrap with span
-        if trace_id and parent_span_id:
-            logger.debug(
-                f"[TemporalTracingResponsesModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}"
-            )
-
-            trace = self._tracer.trace(trace_id)
-
-            async with trace.span(
-                parent_id=parent_span_id,
-                name="model_get_response",
-                input={
-                    "model": str(self.model),
-                    "has_system_instructions": system_instructions is not None,
-                    "input_type": type(input).__name__,
-                    "tools_count": len(tools) if tools else 0,
-                    "handoffs_count": len(handoffs) if handoffs else 0,
-                    "has_output_schema": output_schema is not None,
-                    "model_settings": {
-                        "temperature": model_settings.temperature,
-                        "max_tokens": model_settings.max_tokens,
-                        "reasoning": model_settings.reasoning,
-                    }
-                    if model_settings
-                    else None,
-                },
-            ) as span:
-                try:
-                    # Call the base model
-                    response = await self._base_model.get_response(
-                        system_instructions=system_instructions,
-                        input=input,
-                        model_settings=model_settings,
-                        tools=tools,
-                        output_schema=output_schema,
-                        handoffs=handoffs,
-                        tracing=tracing,
-                        previous_response_id=previous_response_id,
-                        conversation_id=conversation_id,  # type: ignore[call-arg]
-                        prompt=prompt,
-                        **kwargs,
-                    )
-
-                    # Serialize response output items for span tracing
-                    new_items = []
-                    final_output = None
-
-                    if hasattr(response, "output") and response.output:
-                        response_output = response.output if isinstance(response.output, list) else [response.output]
-
-                        for item in response_output:
-                            try:
-                                item_dict = _serialize_item(item)
-                                if item_dict:
-                                    new_items.append(item_dict)
-
-                                    # Extract final_output from message type if available
-                                    if item_dict.get("type") == "message" and not final_output:
-                                        content = item_dict.get("content", [])
-                                        if content and isinstance(content, list):
-                                            for content_part in content:
-                                                if isinstance(content_part, dict) and "text" in content_part:
-                                                    final_output = content_part["text"]
-                                                    break
-                            except Exception as e:
-                                logger.warning(f"Failed to serialize item in temporal tracing model: {e}")
-                                continue
-
-                    # Set span output with structured data
-                    output_data: dict[str, Any] = {
-                        "new_items": new_items,
-                        "final_output": final_output,
-                    }
-                    # Per-call usage for billing; deduped against any turn aggregate
-                    usage_blob = usage_from_openai_response_usage(getattr(response, "usage", None))
-                    if usage_blob:
-                        output_data["usage"] = usage_blob
-                    span.output = output_data  # type: ignore[attr-defined]
-
-                    return response
-
-                except Exception as e:
-                    # Record error in span
-                    span.error = str(e)  # type: ignore[attr-defined]
-                    raise
-        else:
-            # No tracing context, just pass through
-            logger.debug("[TemporalTracingResponsesModel] No tracing context available, calling base model directly")
-            return await self._base_model.get_response(
-                system_instructions=system_instructions,
-                input=input,
-                model_settings=model_settings,
-                tools=tools,
-                output_schema=output_schema,
-                handoffs=handoffs,
-                tracing=tracing,
-                previous_response_id=previous_response_id,
-                conversation_id=conversation_id,  # type: ignore[call-arg]
-                prompt=prompt,
-                **kwargs,
-            )
-
-    @override
-    def stream_response(self, *args, **kwargs):
-        """Streaming is handled via get_response in Temporal activities.
-        Required so the class is concrete and instantiable at runtime."""
-        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
-
-
-class TemporalTracingChatCompletionsModel(Model):
-    """Wrapper for OpenAIChatCompletionsModel that adds AgentEx tracing.
-
-    This is a thin wrapper that adds tracing spans around the base model's
-    get_response() method. It reads tracing context from ContextVars set by
-    the Temporal context interceptor.
-    """
-
-    def __init__(self, base_model: OpenAIChatCompletionsModel, tracer: AsyncTracer):
-        """Initialize the tracing wrapper.
-
-        Args:
-            base_model: The OpenAI ChatCompletions model to wrap
-            tracer: The AgentEx tracer to use
-        """
-        self._base_model = base_model
-        self._tracer = tracer
-        # Expose the model name for compatibility
-        self.model = base_model.model
-
-    @override
-    async def get_response(
-        self,
-        system_instructions: Optional[str],
-        input: Union[str, List[TResponseInputItem]],
-        model_settings: ModelSettings,
-        tools: List[Tool],
-        output_schema: Optional[AgentOutputSchemaBase],
-        handoffs: List[Handoff],
-        tracing: ModelTracing,
-        **kwargs,
-    ) -> ModelResponse:
-        """Get a response from the model with optional tracing.
-
-        If tracing context is available from the interceptor, this wraps the
-        model call in a tracing span. Otherwise, it passes through to the
-        base model without tracing.
-        """
-        # Try to get tracing context from ContextVars
-        task_id = streaming_task_id.get()
-        trace_id = streaming_trace_id.get()
-        parent_span_id = streaming_parent_span_id.get()
-
-        # If we have tracing context, wrap with span
-        if trace_id and parent_span_id:
-            logger.debug(
-                f"[TemporalTracingChatCompletionsModel] Adding tracing span for task_id={task_id}, trace_id={trace_id}"
-            )
-
-            trace = self._tracer.trace(trace_id)
-
-            async with trace.span(
-                parent_id=parent_span_id,
-                name="model_get_response",
-                input={
-                    "model": str(self.model),
-                    "has_system_instructions": system_instructions is not None,
-                    "input_type": type(input).__name__,
-                    "tools_count": len(tools) if tools else 0,
-                    "handoffs_count": len(handoffs) if handoffs else 0,
-                    "has_output_schema": output_schema is not None,
-                    "model_settings": {
-                        "temperature": model_settings.temperature,
-                        "max_tokens": model_settings.max_tokens,
-                    }
-                    if model_settings
-                    else None,
-                },
-            ) as span:
-                try:
-                    # Call the base model
-                    response = await self._base_model.get_response(
-                        system_instructions=system_instructions,
-                        input=input,
-                        model_settings=model_settings,
-                        tools=tools,
-                        output_schema=output_schema,
-                        handoffs=handoffs,
-                        tracing=tracing,
-                        **kwargs,
-                    )
-
-                    # Serialize response output items for span tracing
-                    new_items = []
-                    final_output = None
-
-                    if hasattr(response, "output") and response.output:
-                        response_output = response.output if isinstance(response.output, list) else [response.output]
-
-                        for item in response_output:
-                            try:
-                                item_dict = _serialize_item(item)
-                                if item_dict:
-                                    new_items.append(item_dict)
-
-                                    # Extract final_output from message type if available
-                                    if item_dict.get("type") == "message" and not final_output:
-                                        content = item_dict.get("content", [])
-                                        if content and isinstance(content, list):
-                                            for content_part in content:
-                                                if isinstance(content_part, dict) and "text" in content_part:
-                                                    final_output = content_part["text"]
-                                                    break
-                            except Exception as e:
-                                logger.warning(f"Failed to serialize item in temporal tracing model: {e}")
-                                continue
-
-                    # Set span output with structured data
-                    output_data: dict[str, Any] = {
-                        "new_items": new_items,
-                        "final_output": final_output,
-                    }
-                    # Per-call usage for billing; deduped against any turn aggregate
-                    usage_blob = usage_from_openai_response_usage(getattr(response, "usage", None))
-                    if usage_blob:
-                        output_data["usage"] = usage_blob
-                    span.output = output_data  # type: ignore[attr-defined]
-
-                    return response
-
-                except Exception as e:
-                    # Record error in span
-                    span.error = str(e)  # type: ignore[attr-defined]
-                    raise
-        else:
-            # No tracing context, just pass through
-            logger.debug(
-                "[TemporalTracingChatCompletionsModel] No tracing context available, calling base model directly"
-            )
-            return await self._base_model.get_response(
-                system_instructions=system_instructions,
-                input=input,
-                model_settings=model_settings,
-                tools=tools,
-                output_schema=output_schema,
-                handoffs=handoffs,
-                tracing=tracing,
-                **kwargs,
-            )
-
-    @override
-    def stream_response(self, *args, **kwargs):
-        """Streaming is handled via get_response in Temporal activities.
-        Required so the class is concrete and instantiable at runtime."""
-        raise NotImplementedError("stream_response is not used in Temporal activities - use get_response instead")
diff --git a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
index 282b95e26..bf3dc8006 100644
--- a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
+++ b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py
@@ -1,9 +1,8 @@
-"""Tests that the openai_agents temporal models copy real token usage onto spans.
+"""Tests that the openai_agents streaming model copies real token usage onto spans.
 
 The backend bills per-call usage from ``span.output["usage"]``; these tests
-assert each model writes the framework-reported usage there (and, for the
-streaming model, into the returned ``ModelResponse.usage``) instead of
-dropping it.
+assert the streaming model writes the API-reported usage there (and into the
+returned ``ModelResponse.usage``) instead of dropping it.
 """
 
 from __future__ import annotations
@@ -14,8 +13,7 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
-from agents import ModelResponse, ModelSettings
-from agents.usage import Usage
+from agents import ModelSettings
 from openai.types.responses import Response, ResponseCompletedEvent
 from openai.types.responses.response_usage import (
     ResponseUsage,
@@ -25,10 +23,6 @@
 
 import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm
 from agentex.types.span import Span
-from agentex.lib.core.temporal.plugins.openai_agents.models.temporal_tracing_model import (
-    TemporalTracingResponsesModel,
-    TemporalTracingChatCompletionsModel,
-)
 from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import (
     streaming_task_id,
     streaming_trace_id,
@@ -81,17 +75,6 @@ def tracing_contextvars():
     streaming_parent_span_id.reset(tokens[2])
 
 
-def _agents_usage() -> Usage:
-    return Usage(
-        requests=1,
-        input_tokens=120,
-        output_tokens=80,
-        total_tokens=200,
-        input_tokens_details=InputTokensDetails(cached_tokens=30),
-        output_tokens_details=OutputTokensDetails(reasoning_tokens=40),
-    )
-
-
 EXPECTED_USAGE_BLOB = {
     "input_tokens": 120,
     "output_tokens": 80,
@@ -106,38 +89,6 @@ def _output_dict(span: Span) -> dict[str, Any]:
     return span.output
 
 
-class TestTemporalTracingModels:
-    async def _run_wrapper(self, wrapper_cls) -> Span:
-        tracer = FakeTracer()
-        base_model = MagicMock()
-        base_model.model = "gpt-4o"
-        base_model.get_response = AsyncMock(
-            return_value=ModelResponse(output=[], usage=_agents_usage(), response_id="resp-1")
-        )
-
-        model = wrapper_cls(base_model, tracer)
-        response = await model.get_response(
-            system_instructions=None,
-            input="hello",
-            model_settings=ModelSettings(),
-            tools=[],
-            output_schema=None,
-            handoffs=[],
-            tracing=None,
-        )
-        assert response.usage.input_tokens == 120
-        assert len(tracer.trace_obj.spans) == 1
-        return tracer.trace_obj.spans[0]
-
-    async def test_responses_model_writes_usage_to_span_output(self, tracing_contextvars):
-        span = await self._run_wrapper(TemporalTracingResponsesModel)
-        assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
-
-    async def test_chat_completions_model_writes_usage_to_span_output(self, tracing_contextvars):
-        span = await self._run_wrapper(TemporalTracingChatCompletionsModel)
-        assert _output_dict(span)["usage"] == EXPECTED_USAGE_BLOB
-
-
 class FakeStream:
     def __init__(self, events) -> None:
         self._events = events

From 9e2426d1ba304629cba7c3d76c8a2a5371f2952c Mon Sep 17 00:00:00 2001
From: Levi Lentz <levi.lentz@scale.com>
Date: Thu, 9 Jul 2026 15:18:43 -0700
Subject: [PATCH 9/9] 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.
---
 .../10_async/00_base/030_tracing/README.md    |  10 +-
 src/agentex/lib/adk/_modules/tracing.py       |  71 +++++++----
 .../lib/core/temporal/plugins/__init__.py     |   2 +-
 .../plugins/openai_agents/models/__init__.py  |   2 +-
 src/agentex/lib/core/tracing/usage.py         | 115 ------------------
 tests/lib/adk/test_tracing_module.py          |  45 ++++++-
 tests/lib/core/tracing/test_usage.py          |  90 --------------
 7 files changed, 95 insertions(+), 240 deletions(-)
 delete mode 100644 src/agentex/lib/core/tracing/usage.py
 delete mode 100644 tests/lib/core/tracing/test_usage.py

diff --git a/examples/tutorials/10_async/00_base/030_tracing/README.md b/examples/tutorials/10_async/00_base/030_tracing/README.md
index c0601055e..3c66248e8 100644
--- a/examples/tutorials/10_async/00_base/030_tracing/README.md
+++ b/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
diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py
index be30068a0..7d49bb91c 100644
--- a/src/agentex/lib/adk/_modules/tracing.py
+++ b/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,
diff --git a/src/agentex/lib/core/temporal/plugins/__init__.py b/src/agentex/lib/core/temporal/plugins/__init__.py
index 8377ab5fc..4da2e9ca0 100644
--- a/src/agentex/lib/core/temporal/plugins/__init__.py
+++ b/src/agentex/lib/core/temporal/plugins/__init__.py
@@ -52,4 +52,4 @@
     "streaming_parent_span_id",
     "TemporalStreamingHooks",
     "stream_lifecycle_content",
-]
\ No newline at end of file
+]
diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py
index 6f972a2ed..19dd967d5 100644
--- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py
+++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/__init__.py
@@ -12,4 +12,4 @@
 __all__ = [
     "TemporalStreamingModel",
     "TemporalStreamingModelProvider",
-]
\ No newline at end of file
+]
diff --git a/src/agentex/lib/core/tracing/usage.py b/src/agentex/lib/core/tracing/usage.py
deleted file mode 100644
index 63dbfcd80..000000000
--- a/src/agentex/lib/core/tracing/usage.py
+++ /dev/null
@@ -1,115 +0,0 @@
-"""Helpers for putting LLM token usage onto trace spans in the billable shape.
-
-The AgentEx backend bills token usage from ``application_trace_span`` rows using
-two span fields:
-
-- ``span.data["usage"]`` (+ ``span.data["cost_usd"]``): the per-turn AGGREGATE.
-  Emit at most once per turn, holding that turn's own (per-invocation, not
-  session-cumulative) usage. When a trace contains an aggregate, the backend
-  keeps it and drops all per-call spans in that trace.
-- ``span.output["usage"]``: per-call detail. Summed by the backend, and dropped
-  whenever an aggregate exists in the trace.
-
-Never emit usage on both a rollup span and its per-call children via
-``output["usage"]`` — that double-counts.
-
-Recognized token key spellings (either spelling of a pair works):
-``input_tokens``/``prompt_tokens``, ``output_tokens``/``completion_tokens``,
-``cached_input_tokens``/``cached_tokens``, ``reasoning_tokens``; cost is
-``cost_usd``.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-from collections.abc import Mapping
-
-from agentex.lib.utils.logging import make_logger
-
-logger = make_logger(__name__)
-
-# Key spellings the backend accepts when summing token 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 usage_from_counts(
-    *,
-    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,
-    cost_usd: float | None = None,
-) -> dict[str, Any]:
-    """Build a usage blob with canonical key spellings, omitting None values."""
-    usage: dict[str, Any] = {}
-    if input_tokens is not None:
-        usage["input_tokens"] = input_tokens
-    if output_tokens is not None:
-        usage["output_tokens"] = output_tokens
-    if total_tokens is not None:
-        usage["total_tokens"] = total_tokens
-    if cached_input_tokens is not None:
-        usage["cached_input_tokens"] = cached_input_tokens
-    if reasoning_tokens is not None:
-        usage["reasoning_tokens"] = reasoning_tokens
-    if cost_usd is not None:
-        usage["cost_usd"] = cost_usd
-    return usage
-
-
-def usage_from_openai_response_usage(usage: Any) -> dict[str, Any] | None:
-    """Extract a usage blob from an OpenAI-style usage object.
-
-    Duck-typed so it works for both ``agents.usage.Usage`` (OpenAI Agents SDK)
-    and ``openai.types.responses.ResponseUsage``: reads ``input_tokens``,
-    ``output_tokens``, ``total_tokens``, ``input_tokens_details.cached_tokens``,
-    and ``output_tokens_details.reasoning_tokens``.
-
-    Returns None when there is nothing usable to report.
-    """
-    if usage is None:
-        return None
-
-    input_tokens = getattr(usage, "input_tokens", None)
-    output_tokens = getattr(usage, "output_tokens", None)
-    if input_tokens is None and output_tokens is None:
-        return None
-
-    input_details = getattr(usage, "input_tokens_details", None)
-    output_details = getattr(usage, "output_tokens_details", None)
-    return usage_from_counts(
-        input_tokens=input_tokens,
-        output_tokens=output_tokens,
-        total_tokens=getattr(usage, "total_tokens", None),
-        cached_input_tokens=getattr(input_details, "cached_tokens", None) if input_details is not None else None,
-        reasoning_tokens=getattr(output_details, "reasoning_tokens", None) if output_details is not None else None,
-    )
-
-
-def validate_usage_blob(usage: Mapping[str, Any]) -> dict[str, Any]:
-    """Return the usage mapping as a plain dict, warning on unrecognized shapes.
-
-    The blob is passed through untouched so callers keep full control; the
-    warning catches typos (e.g. ``inputTokens``) that the backend would
-    silently ignore when billing.
-    """
-    blob = dict(usage)
-    if not any(key in RECOGNIZED_USAGE_KEYS for key in blob):
-        logger.warning(
-            "Usage blob has no recognized token keys and will not be billed. "
-            f"Got keys {sorted(blob)}; expected any of {sorted(RECOGNIZED_USAGE_KEYS)}."
-        )
-    return blob
diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py
index adb8641d5..00e4aae65 100644
--- a/tests/lib/adk/test_tracing_module.py
+++ b/tests/lib/adk/test_tracing_module.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"})
diff --git a/tests/lib/core/tracing/test_usage.py b/tests/lib/core/tracing/test_usage.py
deleted file mode 100644
index 110cc07e3..000000000
--- a/tests/lib/core/tracing/test_usage.py
+++ /dev/null
@@ -1,90 +0,0 @@
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-from agentex.lib.core.tracing.usage import (
-    usage_from_counts,
-    validate_usage_blob,
-    usage_from_openai_response_usage,
-)
-
-
-class TestUsageFromCounts:
-    def test_all_fields(self):
-        assert usage_from_counts(
-            input_tokens=100,
-            output_tokens=50,
-            total_tokens=150,
-            cached_input_tokens=20,
-            reasoning_tokens=10,
-            cost_usd=0.0123,
-        ) == {
-            "input_tokens": 100,
-            "output_tokens": 50,
-            "total_tokens": 150,
-            "cached_input_tokens": 20,
-            "reasoning_tokens": 10,
-            "cost_usd": 0.0123,
-        }
-
-    def test_none_values_omitted(self):
-        assert usage_from_counts(input_tokens=100, output_tokens=50) == {
-            "input_tokens": 100,
-            "output_tokens": 50,
-        }
-
-    def test_explicit_zeros_kept(self):
-        assert usage_from_counts(input_tokens=0, output_tokens=0) == {
-            "input_tokens": 0,
-            "output_tokens": 0,
-        }
-
-
-class TestUsageFromOpenAIResponseUsage:
-    def test_none_returns_none(self):
-        assert usage_from_openai_response_usage(None) is None
-
-    def test_object_without_token_fields_returns_none(self):
-        assert usage_from_openai_response_usage(SimpleNamespace(requests=1)) is None
-
-    def test_full_usage_with_details(self):
-        usage = SimpleNamespace(
-            input_tokens=120,
-            output_tokens=80,
-            total_tokens=200,
-            input_tokens_details=SimpleNamespace(cached_tokens=30),
-            output_tokens_details=SimpleNamespace(reasoning_tokens=40),
-        )
-        assert usage_from_openai_response_usage(usage) == {
-            "input_tokens": 120,
-            "output_tokens": 80,
-            "total_tokens": 200,
-            "cached_input_tokens": 30,
-            "reasoning_tokens": 40,
-        }
-
-    def test_usage_without_details(self):
-        usage = SimpleNamespace(
-            input_tokens=10,
-            output_tokens=5,
-            total_tokens=15,
-            input_tokens_details=None,
-            output_tokens_details=None,
-        )
-        assert usage_from_openai_response_usage(usage) == {
-            "input_tokens": 10,
-            "output_tokens": 5,
-            "total_tokens": 15,
-        }
-
-
-class TestValidateUsageBlob:
-    def test_passthrough_recognized_keys(self):
-        blob = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
-        assert validate_usage_blob(blob) == blob
-
-    def test_warns_on_unrecognized_keys(self, caplog):
-        with caplog.at_level("WARNING"):
-            result = validate_usage_blob({"inputTokens": 10})
-        assert result == {"inputTokens": 10}
-        assert any("no recognized token keys" in message for message in caplog.messages)