feat(billing): measure the whole generation pipeline's real token cost

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,
@@ -56,23 +58,35 @@ 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.
         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"``).
         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(
@@ -86,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,

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 ""
 

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.

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"),

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)