From 66c8f11e96f058ef3d3ac801009525d9c70c143c Mon Sep 17 00:00:00 2001
From: guangyu-reflexio <guangyu@reflexio.ai>
Date: Sun, 13 Sep 2026 01:03:40 -0700
Subject: [PATCH 1/2] feat(billing): measure the whole generation pipeline's
 real token cost
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The real provider counts on an `extraction_tokens` event came from one folded
extractor trace, so every stage outside that trace contributed zero — the
deduplicator ("the 2nd LLM call"), the playbook aggregator and reviewer, and the
should-run precheck all burned provider tokens that nothing recorded.

They are now accumulated at the single `litellm.completion` chokepoint.
`_completion_with_hard_timeout` is the only wrapper around it in this module and
has exactly one call site, three lines above `_log_token_usage`, so capturing
there counts every future stage by construction instead of reopening a
call-site campaign each time one is added. The figure is honestly
"generation-pipeline tokens", not "all pipelines": evaluation, optimizer and
search stages run as separate jobs with no generation emit and still contribute
nothing.

Three things this had to get right.

**Never sum the two producers.** In a live run the capture is a superset of the
trace fold — both observe the same completions — so adding them would
double-count every extraction token. The capture replaces it. But only when it
observed a completion: a capture that saw none has no opinion rather than an
answer of zero, so a caller that produced `token_totals` by some other route
still bills its real number instead of silently billing 0. That is why
`RunTokenCapture` counts observations rather than inferring emptiness from
all-zero totals — `claude_code` and `openclaw` genuinely return `Usage(0, 0)`.
`sum_trace_tokens` survives for the durable-resume path, decoded in another
process where the ContextVar is empty by construction.

**Install before the prepare gate, fresh each run.** Before, because the gate
runs the should-run precheck's real `generate_chat_response`. Fresh, because
`_execute_extractor` does not cancel its worker on a timeout: an orphan keeps
writing into the object the previous run left behind, which nothing reads again.
Per-run scope keeps the double-bill guard intact — the characterization test's
"item 2 bills ZERO provider tokens" still holds, and it is what caught an
earlier version of this rule that replaced unconditionally.

**A `ContextVar` default of None, never a shared instance.** A worker submitted
to a bare `ThreadPoolExecutor` never copies the context, so a mutable default
would be one process-global object accumulating across organizations.

Also fixed: `record_extraction_tokens` returned early whenever
`billing_input_tokens <= 0`, discarding real provider cost on exactly the runs
this change exists to capture — a stage with no extraction input text, or the
failure path where `_extraction_input_text` raises and the caller swallows it to
`""`. It now no-ops only when the event would carry nothing at all. This adds no
rows today: an all-zero `extraction_tokens` event maps to an all-zero meter dict
and is dropped as EXEMPT on both planes.

And `window_codec`, which persists `token_totals` as JSON another image may
decode. `decode_plan` splatted it into `RunTokenTotals` kwargs, so the two new
cache fields would have raised `TypeError` inside `window_executor.execute`
mid-rollout, unguarded. Decode now ignores unknown keys (safe for every future
addition) and encode writes only a pinned field set (safe for the image already
deployed, which has no tolerance of its own). Consequence, recorded not
absorbed: the cache split is not carried across a durable resume.

Verified: OSS unit suite 5056 passed, 10 skipped. Four mutations — narrowing the
gate, making the decoder strict, encoding every field, disabling accumulation —
each fail the new tests; restores were checksum-verified.
---
 reflexio/server/billing_meter.py              |  18 ++-
 .../server/llm/_litellm_text_generation.py    |  18 +++
 reflexio/server/llm/token_accounting.py       | 138 +++++++++++++++++-
 .../services/base_generation_service.py       |  41 +++++-
 .../services/durable_learning/window_codec.py |  52 ++++++-
 tests/server/llm/test_litellm_client_unit.py  |  69 +++++++++
 tests/server/llm/test_token_accounting.py     | 115 ++++++++++++++-
 .../test_window_codec_compat.py               | 116 +++++++++++++++
 tests/server/test_billing_meter.py            |  47 ++++++
 9 files changed, 593 insertions(+), 21 deletions(-)
 create mode 100644 tests/server/services/durable_learning/test_window_codec_compat.py

diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py
index 2f885242d..a1d0897c1 100644
--- a/reflexio/server/billing_meter.py
+++ b/reflexio/server/billing_meter.py
@@ -56,10 +56,18 @@ def record_extraction_tokens(
 ) -> None:
     """Emit the Learning cost facet — call only when extraction fired.
 
-    No-op when ``billing_input_tokens <= 0``. Each call mints a fresh
-    ``event_key=f"tok:{uuid4()}"`` so two token emits under the same
-    ``request_id`` (e.g. profile + playbook extraction in one request) never
-    collapse into one billed event downstream.
+    No-op only when the event would carry NOTHING: no metered basis *and* no
+    real provider tokens. Gating on ``billing_input_tokens <= 0`` alone silently
+    discarded real COGS on exactly the runs that motivated capturing it — a
+    stage outside extraction, or a run where ``_extraction_input_text`` raised
+    and the caller swallowed it to ``""`` (``_usage_billing`` around the
+    input-text call). ``count_value`` stays the metered basis, so a zero-basis
+    event carries provider tokens and bills nothing, which is what the two-meter
+    pricing intends.
+
+    Each call mints a fresh ``event_key=f"tok:{uuid4()}"`` so two token emits
+    under the same ``request_id`` (e.g. profile + playbook extraction in one
+    request) never collapse into one billed event downstream.
 
     Args:
         org_id: Organisation identifier.
@@ -72,7 +80,7 @@ def record_extraction_tokens(
         request_id: Optional request correlation ID.
         session_id: Optional session ID.
     """
-    if billing_input_tokens <= 0:
+    if billing_input_tokens <= 0 and prompt_tokens <= 0 and completion_tokens <= 0:
         return
     recorder = record_usage_event_strict if strict else record_usage_event
     recorder(
diff --git a/reflexio/server/llm/_litellm_text_generation.py b/reflexio/server/llm/_litellm_text_generation.py
index 77d96acaa..592222f68 100644
--- a/reflexio/server/llm/_litellm_text_generation.py
+++ b/reflexio/server/llm/_litellm_text_generation.py
@@ -70,6 +70,7 @@
     default_max_tokens_for_model,
     resolve_model_name,
 )
+from reflexio.server.llm.token_accounting import run_token_capture
 
 if TYPE_CHECKING:
     from reflexio.server.llm._litellm_types import LiteLLMConfig
@@ -1194,6 +1195,23 @@ def _log_token_usage(self, params: dict[str, Any], response: Any) -> None:
                 f", cache_write: {cache_creation or 0}, cache_read: {cache_read or 0}"
             )
 
+        # Accumulate into the run-scoped total, if a run installed one. This is
+        # the single chokepoint: `_completion_with_hard_timeout` is the only
+        # wrapper around `litellm.completion` in this module and has exactly one
+        # call site, three lines above the call to this method. Accumulating here
+        # rather than at 19 call sites means every future stage is counted by
+        # construction. `.get()` returns None outside a generation run (and in a
+        # worker whose context was not copied), where contributing nothing is the
+        # correct answer.
+        capture = run_token_capture.get()
+        if capture is not None:
+            capture.observe(
+                prompt_tokens=getattr(usage, "prompt_tokens", None),
+                completion_tokens=getattr(usage, "completion_tokens", None),
+                cache_read_input_tokens=cache_read,
+                cache_write_input_tokens=cache_creation,
+            )
+
         cost = self._compute_cost_usd(response, params.get("model"))
         cost_suffix = f", cost: ${cost:.6f}" if cost is not None else ""
 
diff --git a/reflexio/server/llm/token_accounting.py b/reflexio/server/llm/token_accounting.py
index e029821a5..6299a2eb6 100644
--- a/reflexio/server/llm/token_accounting.py
+++ b/reflexio/server/llm/token_accounting.py
@@ -1,37 +1,159 @@
 """Plain, dependency-free per-run token accounting (OSS-safe).
 
-Folds the per-turn token counts already present on a ToolLoopTrace into a single
-run total. The enterprise billing layer converts this into a TokenUsage; OSS never
-imports reflexio_ext.
+Two producers feed the same ``RunTokenTotals`` shape, and which one is
+authoritative depends on the path:
+
+* **Live runs** use the :data:`run_token_capture` ContextVar. A generation run
+  installs a fresh capture before it starts and every in-process completion adds
+  into it (``_litellm_text_generation._log_token_usage``), so the total covers
+  the whole run — extraction, consolidation/dedup, the playbook aggregator and
+  reviewer, the should-run precheck — not just the extractor's own tool loop.
+* **Durable resume** uses :func:`sum_trace_tokens`, folding the per-turn counts
+  already on a ToolLoopTrace. A resumed window is decoded in a different
+  process, where the ContextVar is empty by construction, so the trace sum is
+  the only value available there.
+
+**They are never summed, and the capture only wins when it observed something.**
+Summing would double-count every extraction token, because in a live run the
+capture is a superset of the trace fold — it sees the extractor's completions
+*and* the ones after it. But a capture that observed ZERO completions has no
+opinion rather than an answer of zero, so the trace fold still stands. That is
+what keeps a caller who supplies ``token_totals`` by some other route from
+silently billing 0, and it is why :class:`RunTokenCapture` counts observations
+instead of inferring emptiness from all-zero totals.
+
+The enterprise billing layer converts this into a TokenUsage; OSS never imports
+reflexio_ext.
 """
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from contextvars import ContextVar
+from dataclasses import dataclass, field
 from typing import Any
 
 
 @dataclass(slots=True)
 class RunTokenTotals:
-    """Accumulated token counts for a single extraction agent run."""
+    """Accumulated token counts for a single generation run.
+
+    The cache fields are **inclusive sub-buckets of ``prompt_tokens``**, not
+    additional tokens: LiteLLM folds cache-creation and cache-read counts into
+    ``prompt_tokens`` before we ever see them
+    (``litellm/llms/anthropic/chat/transformation.py``), matching OpenTelemetry's
+    rule that ``gen_ai.usage.cache_read.input_tokens`` SHOULD be included in
+    ``gen_ai.usage.input_tokens``. They are carried separately only because the
+    provider prices them differently — Anthropic charges cache-read at 0.1x and
+    cache-write at 1.25x base input, a 12.5x spread hidden inside one number.
+
+    **Never add a cache field to ``prompt_tokens``**, and never sum the four
+    fields into a "total": both double-count.
+    """
 
     prompt_tokens: int = 0
     completion_tokens: int = 0
+    cache_read_input_tokens: int = 0
+    cache_write_input_tokens: int = 0
 
-    def add(self, *, prompt_tokens: int | None, completion_tokens: int | None) -> None:
+    def add(
+        self,
+        *,
+        prompt_tokens: int | None,
+        completion_tokens: int | None,
+        cache_read_input_tokens: int | None = None,
+        cache_write_input_tokens: int | None = None,
+    ) -> None:
         """Accumulate token counts, treating None as 0.
 
         Args:
-            prompt_tokens: Prompt token count for one turn, or None if unavailable.
-            completion_tokens: Completion token count for one turn, or None if unavailable.
+            prompt_tokens: Prompt token count for one completion, or None.
+            completion_tokens: Completion token count for one completion, or None.
+            cache_read_input_tokens: Cached-prompt tokens read for this
+                completion, or None. A sub-bucket of ``prompt_tokens``.
+            cache_write_input_tokens: Prompt tokens written to the cache for this
+                completion, or None. A sub-bucket of ``prompt_tokens``.
         """
         self.prompt_tokens += int(prompt_tokens or 0)
         self.completion_tokens += int(completion_tokens or 0)
+        self.cache_read_input_tokens += int(cache_read_input_tokens or 0)
+        self.cache_write_input_tokens += int(cache_write_input_tokens or 0)
+
+
+@dataclass(slots=True)
+class RunTokenCapture:
+    """A run's accumulator plus how many completions it actually observed.
+
+    ``completions`` is not decoration. A run that observed zero completions is
+    NOT the same as one that observed completions reporting zero tokens — the
+    ``claude_code`` and ``openclaw`` providers really do return ``Usage(0, 0)``.
+    Only the first case means "this capture has no opinion", which is what lets
+    a caller fall back to another producer without ever summing the two.
+    """
+
+    totals: RunTokenTotals = field(default_factory=RunTokenTotals)
+    completions: int = 0
+
+    def observe(
+        self,
+        *,
+        prompt_tokens: int | None,
+        completion_tokens: int | None,
+        cache_read_input_tokens: int | None = None,
+        cache_write_input_tokens: int | None = None,
+    ) -> None:
+        """Record one completion's usage, counting the observation itself."""
+        self.completions += 1
+        self.totals.add(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            cache_read_input_tokens=cache_read_input_tokens,
+            cache_write_input_tokens=cache_write_input_tokens,
+        )
+
+
+#: The run-scoped capture, or None when no run has installed one.
+#:
+#: ``default=None`` is load-bearing: a mutable default would be ONE object
+#: shared by every context that never installs its own — including workers
+#: submitted to a bare ``ThreadPoolExecutor`` that does not copy the context
+#: (``agent_success_evaluation/regen_jobs.py``) — and that object would
+#: accumulate across organizations. A None default makes those sites contribute
+#: nothing, which is the correct answer for work outside a generation run.
+run_token_capture: ContextVar[RunTokenCapture | None] = ContextVar(
+    "reflexio_run_token_capture", default=None
+)
+
+
+def begin_run_token_capture() -> RunTokenCapture:
+    """Install a FRESH capture for the run that is about to start.
+
+    Returns the object so the caller can read it back at the end of the run.
+
+    Replacing the object rather than zeroing the previous one is deliberate.
+    A worker that outlived its parent — ``_execute_extractor`` does not cancel
+    the thread on ``FuturesTimeoutError`` — keeps a reference to the OLD object
+    and can keep adding to it. Because the next run reads a different object,
+    those late writes cannot contaminate it.
+
+    Must be called BEFORE any ``contextvars.copy_context()`` that the run
+    performs: a copied context shares the capture *object*, but not a later
+    ``set()``, so a capture installed inside the worker would not propagate
+    back out.
+
+    Returns:
+        RunTokenCapture: The newly installed, zeroed capture.
+    """
+    capture = RunTokenCapture()
+    run_token_capture.set(capture)
+    return capture
 
 
 def sum_trace_tokens(trace: Any) -> RunTokenTotals:
     """Fold a ToolLoopTrace's per-turn token counts into one RunTokenTotals.
 
+    Used by the durable-resume path only — see the module docstring. The trace
+    carries no cache breakdown, so those fields stay 0.
+
     Args:
         trace: A ToolLoopTrace (or duck-typed equivalent) with a ``turns`` attribute.
             Each turn may have ``prompt_tokens`` and ``completion_tokens`` attributes.
diff --git a/reflexio/server/services/base_generation_service.py b/reflexio/server/services/base_generation_service.py
index 05351fabd..697d4e2ad 100644
--- a/reflexio/server/services/base_generation_service.py
+++ b/reflexio/server/services/base_generation_service.py
@@ -15,7 +15,10 @@
 from reflexio.server.api_endpoints.request_context import RequestContext
 from reflexio.server.llm._litellm_types import ModelProvenance
 from reflexio.server.llm.litellm_client import LiteLLMClient
-from reflexio.server.llm.token_accounting import RunTokenTotals
+from reflexio.server.llm.token_accounting import (
+    RunTokenTotals,
+    begin_run_token_capture,
+)
 from reflexio.server.services.base_generation import (
     BatchProgressMixin,
     ConfigFilterMixin,
@@ -665,6 +668,23 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None:
             return None
 
         generation_start = time.perf_counter()
+        # Install the run-scoped token capture BEFORE the prepare gate, not
+        # beside the `_last_token_totals` reset below. The gate runs the
+        # should-run precheck, which is a real `generate_chat_response` call
+        # (`_should_run.py`); installing any later would leave its cost
+        # invisible, which is the whole class of omission this change exists to
+        # end. It is also, necessarily, before `_execute_extractor`'s
+        # `contextvars.copy_context()` — a copied context shares this object but
+        # not a later `set()`, so a capture installed deeper would never
+        # propagate back out.
+        #
+        # Per-run scope is preserved: `compute_generation` runs once per item,
+        # and each run installs a FRESH object. That is what keeps the
+        # double-bill guard intact (`test_base_generation_service_characterization`:
+        # "item 2 bills ZERO provider tokens") and what makes a worker orphaned
+        # by `_execute_extractor`'s timeout harmless — it keeps writing into the
+        # object the previous run left behind, which nothing reads again.
+        token_capture = begin_run_token_capture()
         prepared = self._prepare_generation_run(request)
         if prepared is None:
             return None
@@ -690,6 +710,25 @@ def compute_generation(self, request: TRequest) -> GenerationComputePlan | None:
             self._mark_extraction_runs_finalization_failed(exc)
             raise
 
+        # The capture now covers everything this run did in-process, not just
+        # the extractor's own tool loop: the should-run precheck in the prepare
+        # gate, and `_resolve_write_plan` above, which runs the deduplicator
+        # ("the 2nd LLM call") plus the playbook aggregator and reviewer. It
+        # REPLACES, and is never
+        # summed with, the trace fold `_execute_extractor` stored a moment ago —
+        # both observe the same completions, so adding them would double-count
+        # every extraction token.
+        #
+        # Only when it actually observed a completion, though. A capture that
+        # saw none has no opinion, not an answer of zero, so whatever the
+        # extractor reported stands: an extractor that produced its totals by
+        # some other route must not silently bill 0. In a live run the capture
+        # always sees the extractor's own completions (they go through the one
+        # `litellm.completion` wrapper), so it wins and brings the later stages
+        # with it. The trace fold also survives for the durable-resume path,
+        # which decodes in another process where this ContextVar is empty.
+        if token_capture.completions:
+            self._last_token_totals = token_capture.totals
         generated_count = self._count_generated_results(result)
         billable_count = (
             self._count_retained_online_learnings(write_plan)
diff --git a/reflexio/server/services/durable_learning/window_codec.py b/reflexio/server/services/durable_learning/window_codec.py
index b5f00c6b1..9897c9a2a 100644
--- a/reflexio/server/services/durable_learning/window_codec.py
+++ b/reflexio/server/services/durable_learning/window_codec.py
@@ -1,6 +1,28 @@
-"""Versioned JSON representation of a computed window; no executable payloads."""
+"""Versioned JSON representation of a computed window; no executable payloads.
+
+**Adding a field to a persisted payload is a cross-version contract change.**
+A window written by a new image can be decoded by an OLD one mid-rollout, and
+``decode_plan`` used to splat ``token_totals`` straight into ``RunTokenTotals``
+kwargs — so a new field raised ``TypeError`` inside ``window_executor.execute``,
+unguarded, on a durable-learning path. Two defences, and they point in opposite
+directions on purpose:
+
+* :func:`decode_plan` now IGNORES keys it does not know, so any *future* field
+  addition is safe to decode.
+* :data:`_PERSISTED_TOKEN_FIELDS` pins what :func:`encode_plan` *writes*, so an
+  old image (which does not have the tolerance above) never sees a key it cannot
+  handle. Tolerance alone could not protect this first addition, because the
+  image that must tolerate it is the one already deployed.
+
+Consequence, stated rather than absorbed: the cache sub-buckets are not carried
+across a durable resume and read back as 0 there. That path's token totals are
+a fallback for the run-scoped accumulator, so the loss is bounded to resumed
+windows. Widen ``_PERSISTED_TOKEN_FIELDS`` once no pre-cache-field image is
+running.
+"""
 
 import time
+from dataclasses import fields as dataclass_fields
 from typing import Any
 
 from pydantic_core import to_jsonable_python
@@ -20,6 +42,30 @@
     ProfileWritePlan,
 )
 
+#: The ``RunTokenTotals`` fields ``encode_plan`` is allowed to persist. See the
+#: module docstring: this is a compatibility floor, not a description of the
+#: dataclass, and it must not be regenerated from it.
+_PERSISTED_TOKEN_FIELDS = ("prompt_tokens", "completion_tokens")
+
+
+def _encode_token_totals(totals: RunTokenTotals | None) -> dict[str, int] | None:
+    """Serialise only the fields every deployed image can decode."""
+    if totals is None:
+        return None
+    return {name: getattr(totals, name) for name in _PERSISTED_TOKEN_FIELDS}
+
+
+def _decode_token_totals(raw: Any) -> RunTokenTotals | None:
+    """Rebuild ``RunTokenTotals``, ignoring keys this build does not know.
+
+    Unknown keys are dropped rather than raising, so a window written by a newer
+    image decodes here instead of failing the whole execution.
+    """
+    if not raw:
+        return None
+    known = {f.name for f in dataclass_fields(RunTokenTotals)}
+    return RunTokenTotals(**{k: v for k, v in raw.items() if k in known})
+
 
 def encode_plan(plan: GenerationComputePlan | None, service: Any) -> dict[str, Any]:
     if plan is None:
@@ -32,7 +78,7 @@ def encode_plan(plan: GenerationComputePlan | None, service: Any) -> dict[str, A
             "generated_count": plan.generated_count,
             "billable_count": plan.billable_count,
             "extraction_run_ids": plan.extraction_run_ids,
-            "token_totals": plan.token_totals,
+            "token_totals": _encode_token_totals(plan.token_totals),
             "model_provenance": service._last_model_provenance,
             "stats": service._last_extractor_run_stats,
             "finalization_result": plan.finalization_result,
@@ -68,7 +114,7 @@ def decode_plan(
                     **raw["consolidation_provenance"]
                 )
             write_plan = PlaybookWritePlan(**raw)
-    totals = RunTokenTotals(**data["token_totals"]) if data["token_totals"] else None
+    totals = _decode_token_totals(data["token_totals"])
     service._last_extraction_run_ids = data["extraction_run_ids"]
     service._last_token_totals = totals
     service._last_model_provenance = (
diff --git a/tests/server/llm/test_litellm_client_unit.py b/tests/server/llm/test_litellm_client_unit.py
index c5fad018b..3b5fefbaf 100644
--- a/tests/server/llm/test_litellm_client_unit.py
+++ b/tests/server/llm/test_litellm_client_unit.py
@@ -62,6 +62,10 @@
     create_litellm_client,
 )
 from reflexio.server.llm.llm_utils import make_strict_json_schema
+from reflexio.server.llm.token_accounting import (
+    begin_run_token_capture,
+    run_token_capture,
+)
 
 # ---------------------------------------------------------------------------
 # Pydantic models used for structured-output tests
@@ -3274,6 +3278,71 @@ def test_with_anthropic_cache_stats(self, client):
         # Should not raise
         client._log_token_usage({"model": "claude-3"}, response)
 
+    @staticmethod
+    def _response(prompt, completion, *, cache_write=None, cache_read=None):
+        response = MagicMock()
+        response.usage.prompt_tokens = prompt
+        response.usage.completion_tokens = completion
+        response.usage.total_tokens = prompt + completion
+        response.usage.prompt_tokens_details = None
+        response.usage.cache_creation_input_tokens = cache_write
+        response.usage.cache_read_input_tokens = cache_read
+        return response
+
+    def test_every_completion_in_a_run_accumulates_into_one_capture(self, client):
+        """The premise of run-scoped capture: a SECOND call still counts.
+
+        This is what makes the figure "generation-pipeline tokens" rather than
+        "extraction tokens". In a real run the later calls are the deduplicator,
+        the playbook aggregator and reviewer, and the should-run precheck — all
+        of which used to contribute nothing because the only producer folded one
+        extractor trace.
+        """
+        capture = begin_run_token_capture()
+        client._log_token_usage({"model": "gpt-4o"}, self._response(100, 10))
+        client._log_token_usage({"model": "gpt-4o"}, self._response(30, 4))
+
+        assert capture.completions == 2
+        assert capture.totals.prompt_tokens == 130
+        assert capture.totals.completion_tokens == 14
+
+    def test_cache_sub_buckets_are_captured_without_inflating_input(self, client):
+        """Cache counts ride alongside prompt_tokens, never added into it.
+
+        LiteLLM has already folded them in; re-adding would overstate the most
+        expensive half of the bill.
+        """
+        capture = begin_run_token_capture()
+        client._log_token_usage(
+            {"model": "claude-3"},
+            self._response(1000, 50, cache_write=100, cache_read=800),
+        )
+
+        assert capture.totals.prompt_tokens == 1000
+        assert capture.totals.cache_write_input_tokens == 100
+        assert capture.totals.cache_read_input_tokens == 800
+
+    def test_no_capture_installed_is_a_no_op(self, client):
+        """Outside a generation run there is nothing to accumulate into.
+
+        The chokepoint is on every completion in the process, including ones
+        with no run around them; it must not need one.
+        """
+        run_token_capture.set(None)
+        client._log_token_usage({"model": "gpt-4o"}, self._response(10, 1))
+        assert run_token_capture.get() is None
+
+    def test_a_response_without_usage_is_not_counted_as_an_observation(self, client):
+        """No usage means nothing was observed — not a zero-token observation.
+
+        The observation count decides whether the capture overrides another
+        producer, so counting a usage-less response would let an empty capture
+        overrule a real total.
+        """
+        capture = begin_run_token_capture()
+        client._log_token_usage({"model": "gpt-4o"}, MagicMock(spec=[]))
+        assert capture.completions == 0
+
 
 # ===================================================================
 # create_litellm_client convenience function tests
diff --git a/tests/server/llm/test_token_accounting.py b/tests/server/llm/test_token_accounting.py
index 1722a3ccf..e43c8ae61 100644
--- a/tests/server/llm/test_token_accounting.py
+++ b/tests/server/llm/test_token_accounting.py
@@ -1,13 +1,23 @@
 """Unit tests for the OSS per-run token accounting helpers.
 
-Covers ``RunTokenTotals.add`` (including the ``int(x or 0)`` None->0 coercion)
-and ``sum_trace_tokens`` (the missing/empty ``turns`` guard and summation across
-turns).
+Covers ``RunTokenTotals.add`` (including the ``int(x or 0)`` None->0 coercion),
+``sum_trace_tokens`` (the missing/empty ``turns`` guard and summation across
+turns), and the run-scoped ``RunTokenCapture`` — whose three load-bearing
+properties are the absent-by-default ContextVar, the fresh object per run, and
+the observation COUNT that distinguishes "saw nothing" from "saw zeros".
 """
 
+import contextvars
+from concurrent.futures import ThreadPoolExecutor
 from types import SimpleNamespace
 
-from reflexio.server.llm.token_accounting import RunTokenTotals, sum_trace_tokens
+from reflexio.server.llm.token_accounting import (
+    RunTokenCapture,
+    RunTokenTotals,
+    begin_run_token_capture,
+    run_token_capture,
+    sum_trace_tokens,
+)
 
 
 def test_run_token_totals_default_zero() -> None:
@@ -105,3 +115,100 @@ def test_sum_trace_tokens_turn_none_token_values() -> None:
     totals = sum_trace_tokens(trace)
     assert totals.prompt_tokens == 8
     assert totals.completion_tokens == 6
+
+
+# ── The run-scoped capture (Change 1) ───────────────────────────────────────
+
+
+def test_cache_buckets_accumulate_separately_from_prompt_tokens() -> None:
+    """The cache fields are carried alongside, never folded into, prompt_tokens.
+
+    LiteLLM has already added cache-creation and cache-read into prompt_tokens
+    before we see them, so adding them again here would double-count the most
+    expensive half of the bill.
+    """
+    totals = RunTokenTotals()
+    totals.add(
+        prompt_tokens=1000,
+        completion_tokens=50,
+        cache_read_input_tokens=800,
+        cache_write_input_tokens=100,
+    )
+    assert totals.prompt_tokens == 1000
+    assert totals.cache_read_input_tokens == 800
+    assert totals.cache_write_input_tokens == 100
+
+
+def test_capture_counts_observations_not_just_tokens() -> None:
+    """A completion reporting 0/0 still counts as an observation.
+
+    ``claude_code`` and ``openclaw`` genuinely return ``Usage(0, 0)``. Inferring
+    "nothing was observed" from all-zero totals would make those runs fall back
+    to another producer, which is the bug this counter exists to prevent.
+    """
+    capture = RunTokenCapture()
+    assert capture.completions == 0
+    capture.observe(prompt_tokens=0, completion_tokens=0)
+    assert capture.completions == 1
+    assert capture.totals.prompt_tokens == 0
+
+
+def test_begin_run_token_capture_installs_a_fresh_object_each_time() -> None:
+    """Each run gets a NEW object, so a leaked writer cannot reach the next run.
+
+    ``_execute_extractor`` does not cancel its worker on a timeout, so an
+    orphaned thread can keep calling ``observe`` after the parent has moved on.
+    It holds the old object; replacing rather than zeroing is what makes those
+    late writes harmless.
+    """
+    first = begin_run_token_capture()
+    first.observe(prompt_tokens=10, completion_tokens=1)
+
+    second = begin_run_token_capture()
+    assert second is not first
+    assert second.completions == 0
+    assert second.totals.prompt_tokens == 0
+
+    # The "orphan" writes into the object it still holds; the live run is unmoved.
+    first.observe(prompt_tokens=999, completion_tokens=999)
+    assert run_token_capture.get() is second
+    assert second.totals.prompt_tokens == 0
+
+
+def test_capture_is_absent_by_default() -> None:
+    """Outside a generation run there is no capture at all — not a shared one.
+
+    A mutable ContextVar default would be one process-global object collecting
+    tokens across organisations, because a worker submitted to a bare
+    ThreadPoolExecutor never copies the context.
+    """
+
+    def read_in_uncopied_worker() -> RunTokenCapture | None:
+        with ThreadPoolExecutor(max_workers=1) as pool:
+            return pool.submit(run_token_capture.get).result()
+
+    begin_run_token_capture()
+    # The worker inherits nothing: a new thread starts from an empty context.
+    assert read_in_uncopied_worker() is None
+
+
+def test_copied_context_shares_the_capture_object() -> None:
+    """The install must precede ``copy_context()`` — this pins why.
+
+    A copied context shares the object, so a worker's ``observe`` is visible to
+    the parent. It does NOT share a later ``set()``, which is why installing the
+    capture inside the worker would silently produce zero.
+    """
+    capture = begin_run_token_capture()
+    ctx = contextvars.copy_context()
+
+    def work() -> None:
+        inner = run_token_capture.get()
+        assert inner is capture
+        inner.observe(prompt_tokens=7, completion_tokens=2)
+
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        pool.submit(ctx.run, work).result()
+
+    assert capture.totals.prompt_tokens == 7
+    assert capture.completions == 1
diff --git a/tests/server/services/durable_learning/test_window_codec_compat.py b/tests/server/services/durable_learning/test_window_codec_compat.py
new file mode 100644
index 000000000..afd2711a1
--- /dev/null
+++ b/tests/server/services/durable_learning/test_window_codec_compat.py
@@ -0,0 +1,116 @@
+"""Cross-version tolerance of the persisted window payload.
+
+``token_totals`` is persisted JSON that a *different image* may decode: during a
+rolling deploy, a window written by a new task can be picked up by an old one.
+``decode_plan`` used to splat it straight into ``RunTokenTotals`` kwargs, so
+adding a field raised ``TypeError`` inside ``window_executor.execute`` — on a
+durable-learning path, unguarded.
+
+Both halves of the fix are pinned here, because each covers a direction the
+other cannot:
+
+* decode ignores keys it does not know → a NEW window decodes on an old-enough
+  build, and every future field addition is safe;
+* encode writes only the pinned field set → the image already deployed, which
+  has no tolerance of its own, never receives a key it cannot handle.
+"""
+
+import time
+
+from reflexio.server.llm.token_accounting import RunTokenTotals
+from reflexio.server.services.base_generation_service import PreparedGenerationRun
+from reflexio.server.services.deferred_learning_plan import GenerationComputePlan
+from reflexio.server.services.durable_learning.window_codec import (
+    _PERSISTED_TOKEN_FIELDS,
+    _decode_token_totals,
+    _encode_token_totals,
+    encode_plan,
+)
+
+
+class _FakeService:
+    _last_model_provenance = None
+    _last_extractor_run_stats = {"total": 1, "failed": 0, "timed_out": 0}
+
+
+def _plan(totals: RunTokenTotals | None) -> GenerationComputePlan:
+    return GenerationComputePlan(
+        prepared=PreparedGenerationRun(None, "fake-extractor", "user1"),
+        generated_count=1,
+        billable_count=1,
+        write_plan=None,
+        bookmark_advance=None,
+        generation_start=time.perf_counter(),
+        extraction_run_ids=[],
+        token_totals=totals,
+    )
+
+
+def test_decode_ignores_a_field_this_build_does_not_know() -> None:
+    """A newer image's key must not fail the whole window execution."""
+    totals = _decode_token_totals(
+        {
+            "prompt_tokens": 12,
+            "completion_tokens": 3,
+            "a_field_from_a_newer_image": 99,
+        }
+    )
+    assert totals is not None
+    assert totals.prompt_tokens == 12
+    assert totals.completion_tokens == 3
+
+
+def test_decode_round_trips_the_fields_it_does_know() -> None:
+    assert _decode_token_totals({"prompt_tokens": 5}) == RunTokenTotals(prompt_tokens=5)
+    assert _decode_token_totals(None) is None
+    assert _decode_token_totals({}) is None
+
+
+def test_encode_writes_only_the_pinned_compatibility_floor() -> None:
+    """The cache fields exist on the dataclass but are NOT persisted yet.
+
+    This is the half that protects the image already running, which has no
+    unknown-key tolerance. Widening ``_PERSISTED_TOKEN_FIELDS`` is a deliberate
+    later step, not something to do by regenerating it from the dataclass.
+    """
+    totals = RunTokenTotals(
+        prompt_tokens=100,
+        completion_tokens=20,
+        cache_read_input_tokens=80,
+        cache_write_input_tokens=10,
+    )
+    encoded = _encode_token_totals(totals)
+    assert encoded == {"prompt_tokens": 100, "completion_tokens": 20}
+    assert set(encoded) == set(_PERSISTED_TOKEN_FIELDS)
+    assert _encode_token_totals(None) is None
+
+
+def test_encode_plan_routes_token_totals_through_the_filter() -> None:
+    """Pin the wiring, not just the helper — an unrouted `encode_plan` would
+    reintroduce the break while every helper test stayed green."""
+    encoded = encode_plan(
+        _plan(
+            RunTokenTotals(
+                prompt_tokens=7,
+                completion_tokens=1,
+                cache_read_input_tokens=5,
+            )
+        ),
+        _FakeService(),
+    )
+    assert encoded["token_totals"] == {"prompt_tokens": 7, "completion_tokens": 1}
+
+
+def test_a_new_windows_payload_survives_this_builds_decoder() -> None:
+    """End to end: the shape a future image writes still decodes here."""
+    future_payload = {
+        "prompt_tokens": 42,
+        "completion_tokens": 9,
+        "cache_read_input_tokens": 30,
+        "cache_write_input_tokens": 4,
+        "some_meter_invented_later": 1,
+    }
+    totals = _decode_token_totals(future_payload)
+    assert totals is not None
+    assert totals.prompt_tokens == 42
+    assert totals.cache_read_input_tokens == 30
diff --git a/tests/server/test_billing_meter.py b/tests/server/test_billing_meter.py
index 416bbbc6e..8877cdceb 100644
--- a/tests/server/test_billing_meter.py
+++ b/tests/server/test_billing_meter.py
@@ -67,6 +67,53 @@ def test_record_extraction_tokens_noop_on_negative():
     hook.assert_not_called()
 
 
+def test_real_tokens_survive_a_zero_metered_basis():
+    """The durability axis: no metered basis must not discard real provider cost.
+
+    ``billing_input_tokens`` is a recount of the extraction INPUT text. It is 0
+    on a stage that has no such text, and on the failure path where
+    ``_extraction_input_text`` raises and the caller swallows it to ``""``. The
+    old ``billing_input_tokens <= 0`` gate returned there, throwing away the
+    provider tokens that motivated capturing them at all — and no accumulation
+    test varies this axis, because accumulation happens long before the gate.
+    """
+    with patch(HOOK) as hook:
+        record_extraction_tokens(
+            org_id="org1",
+            billing_input_tokens=0,
+            prompt_tokens=4200,
+            completion_tokens=310,
+            platform_llm=True,
+            platform_storage=None,
+        )
+    hook.assert_called_once()
+    kwargs = hook.call_args.kwargs
+    assert kwargs["prompt_tokens"] == 4200
+    assert kwargs["completion_tokens"] == 310
+    # The metered basis stays 0 — this event carries cost, and bills nothing.
+    assert kwargs["count_value"] == 0
+    assert kwargs["billing_input_tokens"] == 0
+
+
+def test_record_extraction_tokens_noop_when_the_event_would_carry_nothing():
+    """Still a no-op when there is neither a basis nor any provider tokens.
+
+    Widening the gate must not turn every empty run into an event: an all-zero
+    ``extraction_tokens`` event maps to an all-zero meter dict and is dropped
+    downstream as EXEMPT, so emitting one is pure noise.
+    """
+    with patch(HOOK) as hook:
+        record_extraction_tokens(
+            org_id="org1",
+            billing_input_tokens=0,
+            prompt_tokens=0,
+            completion_tokens=0,
+            platform_llm=True,
+            platform_storage=None,
+        )
+    hook.assert_not_called()
+
+
 def test_record_learnings_generated_uses_count_value():
     with patch(HOOK) as hook:
         record_learnings_generated(

From deedccc5ae59d87e710b73e3c41727ac8a32e034 Mon Sep 17 00:00:00 2001
From: guangyu-reflexio <guangyu@reflexio.ai>
Date: Sun, 13 Sep 2026 01:38:59 -0700
Subject: [PATCH 2/2] feat(billing): carry the cache sub-buckets on the
 usage-event wire
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Change 1 captured `cache_read_input_tokens` / `cache_write_input_tokens` at the
completion chokepoint, but nothing could carry them: `UsageEvent` had no such
fields, so the values died at the emit site. This adds the carrier — two
dataclass fields, two kwargs on each of `record_usage_event` /
`record_usage_event_strict` / `record_extraction_tokens` — and passes them from
both emit paths, the synchronous one and the durable window's billing snapshot.

The durable path reads them with `.get(..., 0)` rather than `[...]`. An effects
blob written before these keys existed is still billable, and a KeyError at bill
time would strand a window whose learnings are already committed.

Every docstring says the same thing in the same words, because it is the
property most likely to be "fixed" by someone later: the cache counts are
INCLUSIVE sub-buckets of `prompt_tokens`, folded in by the provider before we
see them. Never add them to it; never sum the four.
---
 reflexio/server/billing_meter.py                   |  8 ++++++++
 .../services/base_generation/_usage_billing.py     |  2 ++
 .../services/durable_learning/window_executor.py   | 10 ++++++++++
 reflexio/server/usage_metrics.py                   | 14 ++++++++++++++
 4 files changed, 34 insertions(+)

diff --git a/reflexio/server/billing_meter.py b/reflexio/server/billing_meter.py
index a1d0897c1..07c13f63f 100644
--- a/reflexio/server/billing_meter.py
+++ b/reflexio/server/billing_meter.py
@@ -45,6 +45,8 @@ def record_extraction_tokens(
     billing_input_tokens: int,
     prompt_tokens: int,
     completion_tokens: int,
+    cache_read_input_tokens: int = 0,
+    cache_write_input_tokens: int = 0,
     platform_llm: bool | None,
     platform_storage: bool | None,
     pipeline: str | None = None,
@@ -74,6 +76,10 @@ def record_extraction_tokens(
         billing_input_tokens: Input-anchored token count (the metered basis).
         prompt_tokens: Real provider prompt tokens (COGS; not billed to customer).
         completion_tokens: Real provider completion tokens (COGS; not billed).
+        cache_read_input_tokens: Cached prompt tokens read, an INCLUSIVE
+            sub-bucket of ``prompt_tokens`` — never add the two.
+        cache_write_input_tokens: Prompt tokens written to the cache, likewise
+            an inclusive sub-bucket of ``prompt_tokens``.
         platform_llm: True iff the platform supplies the LLM for this org.
         platform_storage: True iff the platform supplies storage; None defers to rollup.
         pipeline: Optional pipeline tag (e.g. ``"profile"``).
@@ -94,6 +100,8 @@ def record_extraction_tokens(
         count_value=billing_input_tokens,
         prompt_tokens=prompt_tokens,
         completion_tokens=completion_tokens,
+        cache_read_input_tokens=cache_read_input_tokens,
+        cache_write_input_tokens=cache_write_input_tokens,
         billing_input_tokens=billing_input_tokens,
         platform_llm=platform_llm,
         platform_storage=platform_storage,
diff --git a/reflexio/server/services/base_generation/_usage_billing.py b/reflexio/server/services/base_generation/_usage_billing.py
index 96276b552..f650f9474 100644
--- a/reflexio/server/services/base_generation/_usage_billing.py
+++ b/reflexio/server/services/base_generation/_usage_billing.py
@@ -252,6 +252,8 @@ def _record_billing_learning_events(
                 billing_input_tokens=billing_input_tokens,
                 prompt_tokens=totals.prompt_tokens,
                 completion_tokens=totals.completion_tokens,
+                cache_read_input_tokens=totals.cache_read_input_tokens,
+                cache_write_input_tokens=totals.cache_write_input_tokens,
                 platform_llm=platform_llm,
                 platform_storage=None,
                 pipeline=ctx.get("pipeline"),
diff --git a/reflexio/server/services/durable_learning/window_executor.py b/reflexio/server/services/durable_learning/window_executor.py
index 11c110e05..15f5137f4 100644
--- a/reflexio/server/services/durable_learning/window_executor.py
+++ b/reflexio/server/services/durable_learning/window_executor.py
@@ -225,6 +225,11 @@ def _billing_snapshot(
             ),
             "prompt_tokens": totals.prompt_tokens,
             "completion_tokens": totals.completion_tokens,
+            # `.get(...)` on the read side, not here: an effects blob written by
+            # an older image has no such key, and a KeyError at bill time would
+            # strand a window that already committed its learnings.
+            "cache_read_input_tokens": totals.cache_read_input_tokens,
+            "cache_write_input_tokens": totals.cache_write_input_tokens,
         }
 
     def _bill(self, window: Window, billing: dict[str, Any]) -> None:
@@ -249,6 +254,11 @@ def _bill(self, window: Window, billing: dict[str, Any]) -> None:
             billing_input_tokens=billing["input_tokens"],
             prompt_tokens=billing["prompt_tokens"],
             completion_tokens=billing["completion_tokens"],
+            # Tolerant reads: an effects blob written before the cache fields
+            # existed is still billable, and a KeyError here would strand a
+            # window whose learnings are already committed.
+            cache_read_input_tokens=billing.get("cache_read_input_tokens", 0),
+            cache_write_input_tokens=billing.get("cache_write_input_tokens", 0),
             platform_llm=billing["platform_llm"],
             platform_storage=None,
             pipeline=window.kind,
diff --git a/reflexio/server/usage_metrics.py b/reflexio/server/usage_metrics.py
index 28ff68053..5dee1129e 100644
--- a/reflexio/server/usage_metrics.py
+++ b/reflexio/server/usage_metrics.py
@@ -38,6 +38,12 @@ class UsageEvent:
     count_value: int = 1
     prompt_tokens: int | None = None
     completion_tokens: int | None = None
+    # Inclusive sub-buckets of `prompt_tokens`, NOT additional tokens: the
+    # provider folds them in before we see them. They ride separately only
+    # because they are priced differently (cache read ~0.1x, cache write ~1.25x
+    # base input). Never add them to `prompt_tokens`, and never sum the four.
+    cache_read_input_tokens: int | None = None
+    cache_write_input_tokens: int | None = None
     billing_input_tokens: int | None = None
     platform_llm: bool | None = None
     platform_storage: bool | None = None
@@ -115,6 +121,8 @@ def record_usage_event(
     count_value: int = 1,
     prompt_tokens: int | None = None,
     completion_tokens: int | None = None,
+    cache_read_input_tokens: int | None = None,
+    cache_write_input_tokens: int | None = None,
     billing_input_tokens: int | None = None,
     platform_llm: bool | None = None,
     platform_storage: bool | None = None,
@@ -151,6 +159,8 @@ def record_usage_event(
             count_value=count_value,
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
+            cache_read_input_tokens=cache_read_input_tokens,
+            cache_write_input_tokens=cache_write_input_tokens,
             billing_input_tokens=billing_input_tokens,
             platform_llm=platform_llm,
             platform_storage=platform_storage,
@@ -188,6 +198,8 @@ def record_usage_event_strict(
     count_value: int = 1,
     prompt_tokens: int | None = None,
     completion_tokens: int | None = None,
+    cache_read_input_tokens: int | None = None,
+    cache_write_input_tokens: int | None = None,
     billing_input_tokens: int | None = None,
     platform_llm: bool | None = None,
     platform_storage: bool | None = None,
@@ -227,6 +239,8 @@ def record_usage_event_strict(
             count_value=count_value,
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
+            cache_read_input_tokens=cache_read_input_tokens,
+            cache_write_input_tokens=cache_write_input_tokens,
             billing_input_tokens=billing_input_tokens,
             platform_llm=platform_llm,
             platform_storage=platform_storage,