feat(traces): Show cached prompt tokens in token breakdowns

api/oss/src/core/tracing/utils/trees.py

@@ -318,98 +318,43 @@ def _set_cumulative(span: OTelFlatSpan, costs: dict):
     )
 
 
+TOKEN_KEYS = ("prompt", "completion", "total", "cache_read", "cache_creation")
+
+
 def cumulate_tokens(
     spans_id_tree: OrderedDict,
     spans_idx: Dict[str, OTelFlatSpan],
 ) -> None:
-    def _get_incremental(span: OTelFlatSpan):
-        _tokens = {
-            "prompt": 0.0,
-            "completion": 0.0,
-            "total": 0.0,
-        }
+    def _get(span: OTelFlatSpan, scope: str):
+        attr: dict = span.attributes or {}
 
-        if span.attributes is None:
-            return _tokens
+        tokens = attr.get("ag", {}).get("metrics", {}).get("tokens", {}).get(scope, {})
+        if not isinstance(tokens, dict):
+            tokens = {}
 
-        attr: dict = span.attributes
+        return {key: tokens.get(key, 0.0) for key in TOKEN_KEYS}
 
-        return {
-            "prompt": (
-                attr.get("ag", {})
-                .get("metrics", {})
-                .get("tokens", {})
-                .get("incremental", {})
-                .get("prompt", 0.0)
-            ),
-            "completion": (
-                attr.get("ag", {})
-                .get("metrics", {})
-                .get("tokens", {})
-                .get("incremental", {})
-                .get("completion", 0.0)
-            ),
-            "total": (
-                attr.get("ag", {})
-                .get("metrics", {})
-                .get("tokens", {})
-                .get("incremental", {})
-                .get("total", 0.0)
-            ),
-        }
+    def _get_incremental(span: OTelFlatSpan):
+        return _get(span, "incremental")
 
     def _get_cumulative(span: OTelFlatSpan):
-        _tokens = {
-            "prompt": 0.0,
-            "completion": 0.0,
-            "total": 0.0,
-        }
-
-        if span.attributes is None:
-            return _tokens
-
-        attr: dict = span.attributes
-
-        return {
-            "prompt": (
-                attr.get("ag", {})
-                .get("metrics", {})
-                .get("tokens", {})
-                .get("cumulative", {})
-                .get("prompt", 0.0)
-            ),
-            "completion": (
-                attr.get("ag", {})
-                .get("metrics", {})
-                .get("tokens", {})
-                .get("cumulative", {})
-                .get("completion", 0.0)
-            ),
-            "total": (
-                attr.get("ag", {})
-                .get("metrics", {})
-                .get("tokens", {})
-                .get("cumulative", {})
-                .get("total", 0.0)
-            ),
-        }
+        return _get(span, "cumulative")
 
     def _accumulate(a: dict, b: dict):
-        return {
-            "prompt": a.get("prompt", 0.0) + b.get("prompt", 0.0),
-            "completion": a.get("completion", 0.0) + b.get("completion", 0.0),
-            "total": a.get("total", 0.0) + b.get("total", 0.0),
-        }
+        return {key: a.get(key, 0.0) + b.get(key, 0.0) for key in TOKEN_KEYS}
 
     def _set_cumulative(span: OTelFlatSpan, tokens: dict):
         if span.attributes is None:
             span.attributes = {}
 
-        if (
-            tokens.get("prompt", 0.0) != 0.0
-            or tokens.get("completion", 0.0) != 0.0
-            or tokens.get("total", 0.0) != 0.0
-        ):
+        # Keep the stored shape stable for non-cache traces: cache keys only when non-zero.
+        tokens = {
+            key: value
+            for key, value in tokens.items()
+            if value != 0.0 or key in ("prompt", "completion", "total")
+        }
+
+        if any(value != 0.0 for value in tokens.values()):
             if "ag" not in span.attributes or not isinstance(
                 span.attributes["ag"],
                 dict,

api/oss/tests/pytest/unit/tracing/utils/test_trees.py

@@ -141,6 +141,42 @@ def test_cumulate_tokens_and_costs_propagate_from_children_to_parent():
     assert root_costs["total"] == pytest.approx(1.2)
 
 
+def test_cumulate_tokens_propagates_cache_split_when_present():
+    root = _span(
+        span_id=ROOT_UUID,
+        span_name="root",
+        prompt_tokens=1,
+        completion_tokens=2,
+    )
+    child = _span(
+        span_id=CHILD_A_UUID,
+        parent_id=ROOT_UUID,
+        span_name="child",
+        prompt_tokens=4,
+        completion_tokens=5,
+        start_offset_s=1,
+    )
+    child.attributes["ag"]["metrics"]["tokens"]["incremental"].update(
+        {"cache_read": 100.0, "cache_creation": 20.0}
+    )
+
+    span_idx = parse_span_dtos_to_span_idx([root, child])
+    tree = parse_span_idx_to_span_id_tree(span_idx)
+
+    cumulate_tokens(tree, span_idx)
+
+    root_tokens = span_idx[ROOT_UUID].attributes["ag"]["metrics"]["tokens"][
+        "cumulative"
+    ]
+    assert root_tokens == {
+        "prompt": 5.0,
+        "completion": 7.0,
+        "total": 12.0,
+        "cache_read": 100.0,
+        "cache_creation": 20.0,
+    }
+
+
 def test_cumulate_errors_propagates_scalar_counts_from_children_to_parent():
     root = _span(
         span_id=ROOT_UUID,

sdks/python/agenta/sdk/agents/adapters/vercel/stream.py

@@ -11,6 +11,7 @@
 from ...dtos import AgentResult
 from ...streaming import AgentStream
 from ...utils.wire import sanitize_runner_error
+from ...wire_models import WireAgentUsage
 from .messages import TOOL_APPROVAL_REQUEST
 
 log = get_module_logger(__name__)
@@ -778,12 +779,14 @@ def _approval_tool_call_id(payload: Dict[str, Any]) -> Optional[Any]:
     return None
 
 
+# The wire model is the single source of truth for the usage field set (camelCase on the wire).
+_USAGE_KEYS = tuple(
+    field.alias or name for name, field in WireAgentUsage.model_fields.items()
+)
+
+
 def _usage_metadata(data: Dict[str, Any]) -> Dict[str, Any]:
-    return {
-        key: data[key]
-        for key in ("input", "output", "total", "cost")
-        if data.get(key) is not None
-    }
+    return {key: data[key] for key in _USAGE_KEYS if data.get(key) is not None}
 
 
 def _committed_revision_data(tool_name: Any, output: Any) -> Optional[Dict[str, Any]]:

sdks/python/agenta/sdk/agents/tracing.py

@@ -229,6 +229,14 @@ def record_usage(usage: Optional[Dict[str, Any]]) -> None:
         span.set_attribute("gen_ai.usage.prompt_tokens", input_tokens)
         span.set_attribute("gen_ai.usage.completion_tokens", output_tokens)
         span.set_attribute("gen_ai.usage.total_tokens", int(usage.get("total") or 0))
+        cache_read = usage.get("cacheRead")
+        if cache_read:
+            span.set_attribute("gen_ai.usage.cache_read.input_tokens", int(cache_read))
+        cache_write = usage.get("cacheWrite")
+        if cache_write:
+            span.set_attribute(
+                "gen_ai.usage.cache_creation.input_tokens", int(cache_write)
+            )
         cost = usage.get("cost")
         if cost:
             span.set_attribute("gen_ai.usage.cost", float(cost))

sdks/python/agenta/sdk/agents/wire_models.py

@@ -354,12 +354,18 @@ class WireHarnessCapabilities(_WireModel):
 
 
 class WireAgentUsage(_WireModel):
-    """Token / cost usage rolled onto a workflow span."""
+    """Token / cost usage rolled onto a workflow span.
+
+    ``input`` counts non-cached prompt tokens only; the cached portion is reported
+    separately so consumers can show a breakdown that sums to ``total``.
+    """
 
     input: Optional[int] = None
     output: Optional[int] = None
     total: Optional[int] = None
     cost: Optional[float] = None
+    cache_read: Optional[int] = Field(default=None, alias="cacheRead")
+    cache_write: Optional[int] = Field(default=None, alias="cacheWrite")
 
 
 # ---------------------------------------------------------------------------

sdks/python/oss/tests/pytest/unit/agents/golden/run_result.ok.json

@@ -6,11 +6,11 @@
   ],
   "events": [
     {"type": "message", "text": "Hello!"},
-    {"type": "usage", "input": 10, "output": 5, "total": 15, "cost": 0.001},
+    {"type": "usage", "input": 10, "output": 5, "total": 135, "cost": 0.001, "cacheRead": 100, "cacheWrite": 20},
     {"type": "done", "stopReason": "end_turn"},
     {"text": "an event with no type, dropped on parse"}
   ],
-  "usage": {"input": 10, "output": 5, "total": 15, "cost": 0.001},
+  "usage": {"input": 10, "output": 5, "total": 135, "cost": 0.001, "cacheRead": 100, "cacheWrite": 20},
   "stopReason": "end_turn",
   "capabilities": {
     "textMessages": true,

sdks/python/oss/tests/pytest/unit/agents/test_wire_contract.py

@@ -566,7 +566,15 @@ def test_result_from_wire_parses_ok(golden):
     # The event with no `type` is dropped on parse; the other three survive.
     assert [e.type for e in result.events] == ["message", "usage", "done"]
     assert result.events[0].data == {"type": "message", "text": "Hello!"}
-    assert result.usage == {"input": 10, "output": 5, "total": 15, "cost": 0.001}
+    # input excludes cached tokens — see WireAgentUsage in wire_models.py.
+    assert result.usage == {
+        "input": 10,
+        "output": 5,
+        "total": 135,
+        "cost": 0.001,
+        "cacheRead": 100,
+        "cacheWrite": 20,
+    }
     assert result.stop_reason == "end_turn"
     assert result.session_id == "sess-42"
     assert result.model == "gpt-5.5"

services/runner/src/engines/sandbox_agent/usage.ts

@@ -33,10 +33,22 @@ export function mergePromptAndStreamUsage(
   const promptUsage = promptResult?.usage;
   const inputTokens = promptUsage?.inputTokens ?? streamUsage?.input ?? 0;
   const outputTokens = promptUsage?.outputTokens ?? streamUsage?.output ?? 0;
-  const total = inputTokens + outputTokens || streamUsage?.total || 0;
+  const cacheRead = promptUsage?.cachedReadTokens ?? streamUsage?.cacheRead;
+  const cacheWrite = promptUsage?.cachedWriteTokens ?? streamUsage?.cacheWrite;
+  const total =
+    inputTokens + outputTokens + (cacheRead ?? 0) + (cacheWrite ?? 0) ||
+    streamUsage?.total ||
+    0;
   const cost = streamUsage?.cost ?? 0;
   return total > 0 || cost > 0
-    ? { input: inputTokens, output: outputTokens, total, cost }
+    ? {
+        input: inputTokens,
+        output: outputTokens,
+        total,
+        cost,
+        ...(cacheRead != null ? { cacheRead } : {}),
+        ...(cacheWrite != null ? { cacheWrite } : {}),
+      }
     : undefined;
 }
 

services/runner/src/protocol.ts

@@ -367,19 +367,27 @@ export type AgentEvent =
       output?: number;
       total?: number;
       cost?: number;
+      cacheRead?: number;
+      cacheWrite?: number;
     }
   | { type: "error"; message: string }
   | { type: "done"; stopReason?: string };
 
 /** A live event sink the engines call as each event is built. */
 export type EmitEvent = (event: AgentEvent) => void;
 
-/** Run token/cost totals, rolled up onto the caller's workflow span. */
+/** Run token/cost totals, rolled up onto the caller's workflow span.
+ * `input` counts non-cached prompt tokens only (Pi semantics); the cached portion is
+ * reported separately so consumers can show a breakdown that sums to `total`. */
 export interface AgentUsage {
   input: number;
   output: number;
   total: number;
   cost: number;
+  /** Prompt tokens served from the provider's cache. Absent when the harness doesn't report it. */
+  cacheRead?: number;
+  /** Prompt tokens written to the provider's cache. Absent when the harness doesn't report it. */
+  cacheWrite?: number;
 }
 
 export interface AgentRunRequest {