fix(billing): capture provider usage before the response body is read

reflexio/server/llm/_litellm_text_generation.py

@@ -1208,7 +1208,21 @@ def _log_token_usage(self, params: dict[str, Any], response: Any) -> None:
             capture.observe(
                 prompt_tokens=getattr(usage, "prompt_tokens", None),
                 completion_tokens=getattr(usage, "completion_tokens", None),
-                cache_read_input_tokens=cache_read,
+                # Two providers report the same quantity in two places. Anthropic
+                # puts cache reads at the top level; OpenAI nests them under
+                # `prompt_tokens_details.cached_tokens` -- which this method
+                # already reads, ten lines up, for the log line. Passing only the
+                # top-level value recorded 0 for every OpenAI call while the log
+                # beside it printed the real count.
+                #
+                # SELECT, never sum: they are the same tokens, and both are
+                # sub-buckets of `prompt_tokens`, so adding would double-count
+                # the most expensive half of the bill.
+                cache_read_input_tokens=(
+                    cache_read
+                    if cache_read is not None
+                    else getattr(details, "cached_tokens", None)
+                ),
                 cache_write_input_tokens=cache_creation,
             )
 
@@ -1427,11 +1441,21 @@ def _call_and_parse(
                     response = self._completion_with_hard_timeout(
                         turn_params, turn_hard_timeout
                     )
+                # BEFORE any read of the response body. This call is what
+                # accumulates the run-scoped provider total, and everything
+                # below it can raise: `response.choices[0]` on an empty or
+                # malformed `choices` throws, `_make_request` catches it as a
+                # `LiteLLMClientError` and may advance to the next fallback rung
+                # -- whose usage IS captured. The run then bills the fallback's
+                # tokens while the provider charged for both.
+                #
+                # Nothing here needs the parsed body: `_log_token_usage` reads
+                # only `response.usage`, and returns early when it is absent.
+                self._log_token_usage(turn_params, response)
                 provenance = self._build_model_provenance(response)
                 message = response.choices[0].message  # type: ignore[reportAttributeAccessIssue]
                 content = message.content
                 finish_reason = response.choices[0].finish_reason  # type: ignore[reportAttributeAccessIssue]
-                self._log_token_usage(turn_params, response)
                 self.logger.info(
                     "event=llm_request_end model=%s timeout=%s has_response_format=%s elapsed_seconds=%.3f success=%s",
                     turn_params.get("model"),

tests/server/llm/test_litellm_client_unit.py

@@ -3322,6 +3322,37 @@ def test_cache_sub_buckets_are_captured_without_inflating_input(self, client):
         assert capture.totals.cache_write_input_tokens == 100
         assert capture.totals.cache_read_input_tokens == 800
 
+    def test_openai_nested_cached_tokens_reach_the_capture(self, client):
+        """OpenAI nests cache reads; Anthropic puts them at the top level.
+
+        `prompt_tokens_details.cached_tokens` is read ten lines above the
+        `observe` call for the LOG line, and was not passed to it. Every OpenAI
+        call therefore recorded `cache_read_input_tokens=0` while the log beside
+        it printed the real number -- the capture disagreeing with its own log.
+        """
+        capture = begin_run_token_capture()
+        response = self._response(1000, 50)  # top-level cache fields are None
+        response.usage.prompt_tokens_details = MagicMock(cached_tokens=800)
+        client._log_token_usage({"model": "gpt-4o"}, response)
+
+        assert capture.totals.cache_read_input_tokens == 800
+        # Still a sub-bucket, never added into the input total.
+        assert capture.totals.prompt_tokens == 1000
+
+    def test_the_two_cache_read_sources_are_selected_between_never_summed(self, client):
+        """Both fields populated -> ONE value, because they are the same tokens.
+
+        Without this, "read the nested one too" invites `cache_read + cached`,
+        which double-counts the most expensive half of the bill. The top-level
+        value wins; 800 + 800 == 1600 is the number this forbids.
+        """
+        capture = begin_run_token_capture()
+        response = self._response(1000, 50, cache_read=800)
+        response.usage.prompt_tokens_details = MagicMock(cached_tokens=800)
+        client._log_token_usage({"model": "gpt-4o"}, response)
+
+        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.
 
@@ -4039,6 +4070,42 @@ def _fake(**params):
         client.generate_chat_response(self._messages())
         assert all("fallbacks" not in p for p in seen)
 
+    def test_a_rung_with_an_unreadable_body_still_contributes_its_tokens(
+        self, monkeypatch
+    ):
+        """The provider charged for it, so the run must count it.
+
+        `_log_token_usage` -- the single chokepoint that accumulates the
+        run-scoped provider total -- used to run AFTER `_build_model_provenance`
+        and `response.choices[0].message`. A response carrying real usage but an
+        empty or malformed `choices` raises on that read, the walk catches it as
+        a transport failure and advances to the next rung, and the first rung's
+        tokens are silently dropped from the run while the bill includes them.
+
+        The primary below returns usage=(10, 5) with `choices = []`; the
+        fallback serves normally with its own usage. Both must land.
+        """
+        client = LiteLLMClient(
+            LiteLLMConfig(model="minimax/MiniMax-M3", fallback_models=["zai/glm-5.2"])
+        )
+
+        def _fake(**params):
+            if params["model"] == "minimax/MiniMax-M3":
+                broken = _make_completion_response("ignored")
+                broken.choices = []  # real usage, unreadable body
+                return broken
+            return _make_completion_response("ok")
+
+        monkeypatch.setattr("litellm.completion", _fake)
+
+        capture = begin_run_token_capture()
+        client.generate_chat_response(self._messages())
+
+        # Two completions, two contributions -- not just the one that parsed.
+        assert capture.completions == 2
+        assert capture.totals.prompt_tokens == 20
+        assert capture.totals.completion_tokens == 10
+
     def test_mixed_ladder_no_longer_raises_and_each_rung_gets_own_transport(
         self, monkeypatch
     ):