feat: better cached token usage logging

application/alembic/versions/0031_token_usage_cache_tokens.py

@@ -0,0 +1,35 @@
+"""0031 token_usage cache tokens — persist the prompt-cache breakdown.
+
+Providers report how much of each prompt was served from the prompt cache
+(``cached_tokens``) and, on newer OpenAI-family models, how much was written
+to it (``cache_write_tokens``). The LLM clients already parsed both and the
+usage layer discarded them, so ``token_usage`` could not show a cache hit
+rate. These two nullable columns carry the breakdown; NULL means the
+provider reported nothing (distinct from 0). ``prompt_tokens`` keeps its
+meaning as the provider's total. Idempotent both ways.
+
+Revision ID: 0031_token_usage_cache_tokens
+Revises: 0030_superseded_messages
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+
+revision: str = "0031_token_usage_cache_tokens"
+down_revision: Union[str, None] = "0030_superseded_messages"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cached_tokens integer;")
+    op.execute(
+        "ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cache_write_tokens integer;"
+    )
+
+
+def downgrade() -> None:
+    op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cache_write_tokens;")
+    op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cached_tokens;")

application/llm/base.py

@@ -710,6 +710,7 @@ class BaseLLM(ABC):
         completion_tokens,
         latency_ms,
         cached_tokens=None,
+        cache_write_tokens=None,
         error=None,
     ):
         # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired
@@ -730,6 +731,8 @@ class BaseLLM(ABC):
         }
         if cached_tokens is not None:
             extra["cached_tokens"] = int(cached_tokens)
+        if cache_write_tokens is not None:
+            extra["cache_write_tokens"] = int(cache_write_tokens)
         if error is not None:
             extra["error_class"] = type(error).__name__
         logging.info("llm_gen_finished", extra=extra)
@@ -757,6 +760,7 @@ class BaseLLM(ABC):
         completion_tokens,
         latency_ms,
         cached_tokens=None,
+        cache_write_tokens=None,
         error=None,
     ):
         # Paired with ``llm_stream_start`` so cost dashboards can sum tokens
@@ -774,6 +778,8 @@ class BaseLLM(ABC):
         }
         if cached_tokens is not None:
             extra["cached_tokens"] = int(cached_tokens)
+        if cache_write_tokens is not None:
+            extra["cache_write_tokens"] = int(cache_write_tokens)
         if error is not None:
             extra["error_class"] = type(error).__name__
         logging.info("llm_stream_finished", extra=extra)

application/llm/openai.py

@@ -1307,17 +1307,36 @@ class OpenAILLM(BaseLLM):
         output_details = getattr(usage, "completion_tokens_details", None)
         try:
             cached = int(getattr(input_details, "cached_tokens", 0) or 0)
+            written = int(getattr(input_details, "cache_write_tokens", 0) or 0)
             reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0)
         except (TypeError, ValueError):
             cached = 0
+            written = 0
             reasoning = 0
-        if cached:
-            result["prompt_tokens_details"] = {"cached_tokens": cached}
+        details = self._prompt_cache_details(cached, written)
+        if details:
+            result["prompt_tokens_details"] = details
         if reasoning:
             result["completion_tokens_details"] = {"reasoning_tokens": reasoning}
         self._last_usage = result
         self._last_usage_claimed = False
 
+    @staticmethod
+    def _prompt_cache_details(cached: int, written: int) -> dict:
+        """Build the ``prompt_tokens_details`` breakdown from the two cache bins.
+
+        Only non-zero bins are included so a provider that never reports
+        caching produces no details (and no misleading zero downstream).
+        ``cache_write_tokens`` is reported by newer OpenAI-family deployments
+        that charge for cache writes.
+        """
+        details = {}
+        if cached:
+            details["cached_tokens"] = cached
+        if written:
+            details["cache_write_tokens"] = written
+        return details
+
     @staticmethod
     def _function_call_ids(response):
         """call_ids of every ``function_call`` item in a Responses output.
@@ -1357,9 +1376,11 @@ class OpenAILLM(BaseLLM):
                 "total_tokens": int(getattr(usage, "total_tokens", 0) or prompt + completion),
             }
             cached = int(getattr(input_details, "cached_tokens", 0) or 0)
+            written = int(getattr(input_details, "cache_write_tokens", 0) or 0)
             reasoning = int(getattr(output_details, "reasoning_tokens", 0) or 0)
-            if cached:
-                result["prompt_tokens_details"] = {"cached_tokens": cached}
+            details = self._prompt_cache_details(cached, written)
+            if details:
+                result["prompt_tokens_details"] = details
             if reasoning:
                 result["completion_tokens_details"] = {"reasoning_tokens": reasoning}
             self._last_usage = result

application/storage/db/models.py

@@ -235,6 +235,12 @@ token_usage_table = Table(
     # Added in ``0015_token_usage_model_id``. Canonical model id (catalog
     # name for built-ins, UUID for BYOM); NULL on un-backfilled rows.
     Column("model_id", Text),
+    # Added in ``0031_token_usage_cache_tokens``. Prompt-cache breakdowns of
+    # ``prompt_tokens`` as reported by the provider (reads; and writes on
+    # newer OpenAI-family models). NULL = not reported, distinct from 0 = no
+    # cache activity, so hit-rate queries stay honest across providers.
+    Column("cached_tokens", Integer),
+    Column("cache_write_tokens", Integer),
 )
 
 user_logs_table = Table(

application/storage/db/repositories/token_usage.py

@@ -35,6 +35,8 @@ class TokenUsageRepository:
         request_id: Optional[str] = None,
         model_id: Optional[str] = None,
         timestamp: Optional[datetime] = None,
+        cached_tokens: Optional[int] = None,
+        cache_write_tokens: Optional[int] = None,
     ) -> None:
         # Attribution guard: the ``token_usage_attribution_chk`` CHECK
         # constraint requires at least one of ``user_id`` / ``api_key``
@@ -60,12 +62,14 @@ class TokenUsageRepository:
                 INSERT INTO token_usage (
                     user_id, api_key, agent_id,
                     prompt_tokens, generated_tokens,
+                    cached_tokens, cache_write_tokens,
                     source, request_id, model_id, timestamp
                 )
                 VALUES (
                     :user_id, :api_key,
                     CAST(:agent_id AS uuid),
                     :prompt_tokens, :generated_tokens,
+                    :cached_tokens, :cache_write_tokens,
                     :source, :request_id, :model_id, COALESCE(:timestamp, now())
                 )
                 """
@@ -76,6 +80,8 @@ class TokenUsageRepository:
                 "agent_id": agent_id_uuid,
                 "prompt_tokens": prompt_tokens,
                 "generated_tokens": generated_tokens,
+                "cached_tokens": cached_tokens,
+                "cache_write_tokens": cache_write_tokens,
                 "source": source,
                 "request_id": request_id,
                 "model_id": model_id,

application/usage.py

@@ -131,6 +131,11 @@ def _persist_call_usage(llm, call_usage):
                 agent_id=str(agent_id) if agent_id else None,
                 prompt_tokens=call_usage["prompt_tokens"],
                 generated_tokens=call_usage["generated_tokens"],
+                # Present only when the provider reported the breakdown;
+                # persisted as NULL otherwise so "unknown" never reads as
+                # "0% cache hits".
+                cached_tokens=call_usage.get("cached_tokens"),
+                cache_write_tokens=call_usage.get("cache_write_tokens"),
                 source=(
                     getattr(llm, "_token_usage_source", None) or "agent_stream"
                 ),
@@ -151,6 +156,12 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in
     ``*_tokens_details`` breakdowns (``cached_tokens``,
     ``reasoning_tokens``) back out of these bins — that would break
     parity with what providers bill.
+
+    The prompt-cache sub-bins ARE carried alongside (``cached_tokens``,
+    ``cache_write_tokens``; Anthropic's ``cache_creation_tokens`` maps to
+    the latter) so persistence and the finish events can chart them. They
+    are added only when the provider reported them. The rest of
+    ``call_usage`` (e.g. ``model``) is preserved rather than replaced.
     """
     reported = getattr(llm, "_last_usage", None)
     if not isinstance(reported, dict):
@@ -174,10 +185,22 @@ def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, in
         # Slotted/immutable LLM stand-ins can't record the claim; this
         # call still gets the provider counts, which is correct for them.
         pass
-    return {
+    merged = {
+        **call_usage,
         "prompt_tokens": int(prompt or 0),
         "generated_tokens": int(completion or 0),
     }
+    details = reported.get("prompt_tokens_details")
+    if isinstance(details, dict):
+        cached = details.get("cached_tokens")
+        written = details.get("cache_write_tokens")
+        if written is None:
+            written = details.get("cache_creation_tokens")
+        if cached is not None:
+            merged["cached_tokens"] = int(cached or 0)
+        if written is not None:
+            merged["cache_write_tokens"] = int(written or 0)
+    return merged
 
 
 def gen_token_usage(func):
@@ -224,6 +247,8 @@ def gen_token_usage(func):
                         prompt_tokens=call_usage["prompt_tokens"],
                         completion_tokens=call_usage["generated_tokens"],
                         latency_ms=int((time.monotonic() - started_at) * 1000),
+                        cached_tokens=call_usage.get("cached_tokens"),
+                        cache_write_tokens=call_usage.get("cache_write_tokens"),
                         error=error,
                     )
                 except Exception:
@@ -272,6 +297,8 @@ def stream_token_usage(func):
                         prompt_tokens=call_usage["prompt_tokens"],
                         completion_tokens=call_usage["generated_tokens"],
                         latency_ms=int((time.monotonic() - started_at) * 1000),
+                        cached_tokens=call_usage.get("cached_tokens"),
+                        cache_write_tokens=call_usage.get("cache_write_tokens"),
                         error=error,
                     )
                 except Exception:

tests/llm/handlers/test_round_usage_persistence.py

@@ -179,8 +179,12 @@ class TestPerRoundUsageRows:
                 list(handler.handle_streaming(agent, first, {}, []))
 
         assert len(rows) == 2
-        assert rows[0] == {"prompt_tokens": 100, "generated_tokens": 10}
-        assert rows[1] == {"prompt_tokens": 200, "generated_tokens": 20}
+        # Provider-reported counts replace the estimates; the rest of the
+        # call record (e.g. ``model``) rides along, hence a subset check
+        # rather than dict equality.
+        counts = lambda row: {k: row[k] for k in ("prompt_tokens", "generated_tokens")}  # noqa: E731
+        assert counts(rows[0]) == {"prompt_tokens": 100, "generated_tokens": 10}
+        assert counts(rows[1]) == {"prompt_tokens": 200, "generated_tokens": 20}
 
 
 class TestPreferProviderUsageClaim:

tests/llm/test_base.py

@@ -2347,3 +2347,38 @@ class TestBuildConversationFromMessagesEmpty:
         result = handler._build_conversation_from_messages(messages)
         assert result is not None
         assert len(result.get("queries", [])) >= 1
+
+
+class TestFinishedLogCacheBins:
+    """``cached_tokens`` / ``cache_write_tokens`` ride on the finish events
+    only when a provider reported them, so dashboards can compute a cache
+    hit rate without a sentinel zero polluting providers that never report."""
+
+    def test_stream_finished_carries_cache_bins_when_given(self, caplog):
+        import logging as _logging
+
+        llm = StubLLM()
+        with caplog.at_level(_logging.INFO, logger="root"):
+            llm._emit_stream_finished_log(
+                "m1",
+                prompt_tokens=100,
+                completion_tokens=5,
+                latency_ms=10,
+                cached_tokens=80,
+                cache_write_tokens=7,
+            )
+        evt = next(r for r in caplog.records if r.message == "llm_stream_finished")
+        assert evt.cached_tokens == 80
+        assert evt.cache_write_tokens == 7
+
+    def test_gen_finished_omits_cache_bins_when_none(self, caplog):
+        import logging as _logging
+
+        llm = StubLLM()
+        with caplog.at_level(_logging.INFO, logger="root"):
+            llm._emit_gen_finished_log(
+                "m1", prompt_tokens=100, completion_tokens=5, latency_ms=10
+            )
+        evt = next(r for r in caplog.records if r.message == "llm_gen_finished")
+        assert not hasattr(evt, "cached_tokens")
+        assert not hasattr(evt, "cache_write_tokens")

tests/llm/test_openai_responses.py

@@ -1550,3 +1550,38 @@ def test_responses_gen_stream_retries_without_tools(monkeypatch):
     assert len(create.calls) == 2
     assert "tools" in create.calls[0]
     assert "tools" not in create.calls[1]
+
+
+@pytest.mark.unit
+def test_record_chat_usage_captures_cache_write_bin(monkeypatch):
+    """Newer OpenAI-family deployments report cache writes next to cache reads."""
+    llm = _make_llm(monkeypatch, ModelCapabilities(api_flavor="chat_completions"))
+    llm._record_chat_usage(_ns(
+        prompt_tokens=10,
+        completion_tokens=7,
+        total_tokens=17,
+        prompt_tokens_details=_ns(cached_tokens=4, cache_write_tokens=6),
+        completion_tokens_details=None,
+    ))
+    assert llm._last_usage["prompt_tokens_details"] == {
+        "cached_tokens": 4,
+        "cache_write_tokens": 6,
+    }
+
+
+@pytest.mark.unit
+def test_record_responses_metadata_captures_cache_write_bin(monkeypatch):
+    llm = _make_llm(monkeypatch, _responses_caps())
+    llm._record_responses_metadata(_ns(
+        id="resp_1",
+        usage=_ns(
+            input_tokens=10,
+            output_tokens=7,
+            total_tokens=17,
+            input_tokens_details=_ns(cached_tokens=0, cache_write_tokens=9),
+            output_tokens_details=None,
+        ),
+    ))
+    # A pure cache-write turn (first request of a prefix) has reads=0 and
+    # writes>0; the details must still be recorded.
+    assert llm._last_usage["prompt_tokens_details"] == {"cache_write_tokens": 9}

tests/storage/db/repositories/test_token_usage.py

@@ -290,3 +290,32 @@ class TestBucketedTotals:
         )
         assert len(rows) == 1
         assert rows[0]["prompt_tokens"] == 10
+
+
+class TestCacheBins:
+    def _row(self, pg_conn, user_id):
+        return pg_conn.execute(
+            text(
+                "SELECT cached_tokens, cache_write_tokens FROM token_usage "
+                "WHERE user_id = :u ORDER BY id DESC LIMIT 1"
+            ),
+            {"u": user_id},
+        ).fetchone()
+
+    def test_persists_cache_bins(self, pg_conn):
+        repo = TokenUsageRepository(pg_conn)
+        repo.insert(
+            user_id="cache-user",
+            prompt_tokens=1000,
+            generated_tokens=10,
+            cached_tokens=800,
+            cache_write_tokens=100,
+        )
+        assert tuple(self._row(pg_conn, "cache-user")) == (800, 100)
+
+    def test_cache_bins_default_to_null_not_zero(self, pg_conn):
+        """NULL means "provider did not report"; 0 means "reported no cache
+        activity". Keeping them distinct is what makes a hit-rate query honest."""
+        repo = TokenUsageRepository(pg_conn)
+        repo.insert(user_id="cache-user-2", prompt_tokens=10, generated_tokens=1)
+        assert tuple(self._row(pg_conn, "cache-user-2")) == (None, None)

tests/storage/db/test_migration_0031.py

@@ -0,0 +1,100 @@
+"""Migration round-trip test for 0031_token_usage_cache_tokens."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+
+
+pytestmark = pytest.mark.integration
+
+
+def _alembic_ini() -> Path:
+    return Path(__file__).resolve().parents[3] / "application" / "alembic.ini"
+
+
+def _run_alembic(url: str, *args: str) -> None:
+    subprocess.check_call(
+        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
+        timeout=60,
+        env={**os.environ, "POSTGRES_URI": url},
+    )
+
+
+def _alembic_heads(url: str) -> list[str]:
+    out = subprocess.check_output(
+        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"],
+        timeout=60,
+        env={**os.environ, "POSTGRES_URI": url},
+        text=True,
+    )
+    return [line for line in out.splitlines() if line.strip()]
+
+
+def _alembic_version(conn) -> str:
+    return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()
+
+
+def _column_exists(conn, table: str, column: str) -> bool:
+    row = conn.execute(
+        text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'"
+        ),
+        {"t": table, "c": column},
+    ).fetchone()
+    return row is not None
+
+
+_0031 = "0031_token_usage_cache_tokens"
+_0030 = "0030_superseded_messages"
+
+
+class TestMigration0031RoundTrip:
+    def test_single_head(self, pg_engine):
+        url = pg_engine.url.render_as_string(hide_password=False)
+        assert len(_alembic_heads(url)) == 1
+
+    def test_head_has_cache_columns(self, pg_engine):
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) >= _0031
+            assert _column_exists(conn, "token_usage", "cached_tokens")
+            assert _column_exists(conn, "token_usage", "cache_write_tokens")
+
+    def test_downgrade_drops_then_upgrade_restores(self, pg_engine):
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0030)
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) == _0030
+            assert not _column_exists(conn, "token_usage", "cached_tokens")
+            assert not _column_exists(conn, "token_usage", "cache_write_tokens")
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) >= _0031
+            assert _column_exists(conn, "token_usage", "cached_tokens")
+
+    def test_existing_rows_read_null_cache_bins(self, pg_engine):
+        """Rows written before 0031 must read NULL (unknown), not 0."""
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0030)
+        with pg_engine.begin() as conn:
+            conn.execute(
+                text(
+                    "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens) "
+                    "VALUES ('u-mig31', 10, 1)"
+                )
+            )
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            row = conn.execute(
+                text(
+                    "SELECT cached_tokens, cache_write_tokens FROM token_usage "
+                    "WHERE user_id = 'u-mig31'"
+                )
+            ).fetchone()
+            assert tuple(row) == (None, None)

tests/test_usage.py

@@ -600,3 +600,115 @@ class TestCountPromptTokens:
         ]
         tokens = _count_prompt_tokens(messages, tools=None)
         assert tokens > 0
+
+
+# ── prompt-cache breakdown ────────────────────────────────────────────────────
+#
+# Providers report ``cached_tokens`` (and, on newer OpenAI-family models,
+# ``cache_write_tokens``) as breakdowns of ``prompt_tokens``. The prompt bin
+# stays the provider total (never subtract the details back out); the two
+# sub-bins ride alongside so persistence and the finish events can chart a
+# cache hit rate.
+
+
+class _ReportingLLM:
+    decoded_token = {"sub": "user_1"}
+    user_api_key = None
+    agent_id = None
+    token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
+
+    def __init__(self, details):
+        self._last_usage = {
+            "prompt_tokens": 1000,
+            "completion_tokens": 10,
+            "total_tokens": 1010,
+        }
+        if details is not None:
+            self._last_usage["prompt_tokens_details"] = details
+        self._last_usage_claimed = False
+        self.emitted = []
+
+    def _emit_gen_finished_log(self, model, **kwargs):
+        self.emitted.append(kwargs)
+
+
+@pytest.mark.unit
+def test_prefer_provider_usage_carries_cache_bins():
+    from application.usage import _prefer_provider_usage
+
+    llm = _ReportingLLM({"cached_tokens": 800, "cache_write_tokens": 100})
+    usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1})
+    assert usage["prompt_tokens"] == 1000
+    assert usage["generated_tokens"] == 10
+    assert usage["cached_tokens"] == 800
+    assert usage["cache_write_tokens"] == 100
+
+
+@pytest.mark.unit
+def test_prefer_provider_usage_maps_anthropic_cache_creation_to_writes():
+    from application.usage import _prefer_provider_usage
+
+    llm = _ReportingLLM({"cached_tokens": 3, "cache_creation_tokens": 9})
+    usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1})
+    assert usage["cached_tokens"] == 3
+    assert usage["cache_write_tokens"] == 9
+
+
+@pytest.mark.unit
+def test_prefer_provider_usage_without_details_has_no_cache_keys():
+    from application.usage import _prefer_provider_usage
+
+    llm = _ReportingLLM(None)
+    usage = _prefer_provider_usage(llm, {"prompt_tokens": 1, "generated_tokens": 1})
+    assert "cached_tokens" not in usage
+    assert "cache_write_tokens" not in usage
+
+
+@pytest.mark.unit
+def test_prefer_provider_usage_keeps_the_rest_of_the_call_record():
+    from application.usage import _prefer_provider_usage
+
+    llm = _ReportingLLM(None)
+    usage = _prefer_provider_usage(
+        llm, {"prompt_tokens": 1, "generated_tokens": 1, "model": "m"}
+    )
+    assert usage["model"] == "m"
+
+
+@pytest.mark.unit
+def test_decorator_persists_and_emits_cache_bins(monkeypatch):
+    _install_fake_token_repo(monkeypatch)
+    llm = _ReportingLLM({"cached_tokens": 800, "cache_write_tokens": 100})
+
+    @gen_token_usage
+    def wrapped(self, model, messages, stream, tools, **kwargs):
+        _ = (model, messages, stream, tools, kwargs)
+        return "ok"
+
+    wrapped(llm, "m", [{"role": "user", "content": "hi"}], False, None)
+
+    row = _FakeTokenUsageRepo.last_instance.inserted[0]
+    assert row["prompt_tokens"] == 1000
+    assert row["cached_tokens"] == 800
+    assert row["cache_write_tokens"] == 100
+    assert llm.emitted[0]["cached_tokens"] == 800
+    assert llm.emitted[0]["cache_write_tokens"] == 100
+
+
+@pytest.mark.unit
+def test_decorator_omits_cache_bins_when_provider_reports_none(monkeypatch):
+    _install_fake_token_repo(monkeypatch)
+    llm = _ReportingLLM(None)
+
+    @gen_token_usage
+    def wrapped(self, model, messages, stream, tools, **kwargs):
+        _ = (model, messages, stream, tools, kwargs)
+        return "ok"
+
+    wrapped(llm, "m", [{"role": "user", "content": "hi"}], False, None)
+
+    row = _FakeTokenUsageRepo.last_instance.inserted[0]
+    assert row["cached_tokens"] is None
+    assert row["cache_write_tokens"] is None
+    assert llm.emitted[0]["cached_tokens"] is None
+    assert llm.emitted[0]["cache_write_tokens"] is None