From 9ecfda613eb817223d2ebeb6b6f7dff2e4fc7a22 Mon Sep 17 00:00:00 2001 From: guangyu-reflexio Date: Mon, 14 Sep 2026 16:02:41 -0700 Subject: [PATCH] fix(billing): capture provider usage before the response body is read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two CodeRabbit findings on #503, both verified against the merged code before being taken. #503 was merged without their being addressed; this is the follow-up. ## A rung whose body will not parse still cost money `_log_token_usage` is the single chokepoint that accumulates the run-scoped provider total. It ran AFTER `_build_model_provenance(response)` and `response.choices[0].message` — so a response carrying real usage with an empty or malformed `choices` raised on that read before the accumulation happened. The walk catches it as a transport failure and advances to the next fallback rung, whose usage IS captured, and the run then reports the fallback's tokens while the provider charged for both. The call moves to immediately after `_completion_with_hard_timeout` returns. Nothing it needs comes from the parsed body: it reads only `response.usage` and returns early when that is absent. The mutation test shows the mechanism verbatim: ERROR event=llm_request_end model=minimax/MiniMax-M3 success=False error_type=IndexError error=list index out of range INFO event=llm_fallback_used primary_model=minimax/MiniMax-M3 served_model=zai/glm-5.2 reason=transport_error — two completions billed, one counted. ## The capture disagreed with its own log line Anthropic reports cache reads at the top level; OpenAI nests them under `prompt_tokens_details.cached_tokens`. `_log_token_usage` already read the nested field — ten lines above the `observe` call, for the log string — and passed only the top-level one to the capture. Every OpenAI call therefore recorded `cache_read_input_tokens=0` beside a log line printing the real number. SELECT between the two, 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. A third test pins that specifically — the forbidden answer is 1600 where both fields read 800. ## Verified - Three mutations, each caught, each restored and checksum-verified with `shasum -a 256 -c`: restoring the old call order, passing only the top-level cache field, and summing the two cache sources. - OSS unit suite: 5059 passed, 10 skipped, 6 subtests passed. - `uv run ruff check reflexio/ tests/` — clean; `ruff format` applied. - `uv run pyright reflexio/server/llm` — 0 errors, 0 warnings. --- .../server/llm/_litellm_text_generation.py | 28 +++++++- tests/server/llm/test_litellm_client_unit.py | 67 +++++++++++++++++++ 2 files changed, 93 insertions(+), 2 deletions(-) diff --git a/reflexio/server/llm/_litellm_text_generation.py b/reflexio/server/llm/_litellm_text_generation.py index 592222f6..1d978426 100644 --- a/reflexio/server/llm/_litellm_text_generation.py +++ b/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"), diff --git a/tests/server/llm/test_litellm_client_unit.py b/tests/server/llm/test_litellm_client_unit.py index 3b5fefba..85903604 100644 --- a/tests/server/llm/test_litellm_client_unit.py +++ b/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 ):