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

[guangyu-reflexio]

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

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.

Summary by CodeRabbit

• Bug Fixes
  • Improved token usage tracking for cached tokens across supported providers without double-counting.
  • Token usage is now recorded even when response content is malformed or cannot be fully processed.
  • Fallback responses retain accurate primary and fallback token totals.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Essentials

Run ID: bff0b342-d66d-40c2-b654-e6793ee30b95

📥 Commits

Reviewing files that changed from the base of the PR and between ef30e58 and 9ecfda6.

📒 Files selected for processing (2)

• reflexio/server/llm/_litellm_text_generation.py
• tests/server/llm/test_litellm_client_unit.py

Included review availability: Your plan provides up to 5 included reviews per hour; 1 remains after this review.

---

📝 Walkthrough

Walkthrough

The LLM client now captures nested cached-token counts without double-counting and logs token usage before response parsing. Tests cover cache-token precedence and fallback responses with unreadable primary bodies.

Changes

Token usage capture

Layer / File(s)	Summary
Cache token selection and validation reflexio/server/llm/_litellm_text_generation.py, tests/server/llm/test_litellm_client_unit.py	Token capture uses nested OpenAI cached-token counts when the top-level cache-read value is absent. A top-level value takes precedence without summing duplicate counts.
Early usage logging and fallback validation reflexio/server/llm/_litellm_text_generation.py, tests/server/llm/test_litellm_client_unit.py	_log_token_usage runs immediately after completion and before response parsing. Tests verify that usage from an unreadable primary response is retained before fallback totals are added.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: yyiilluu

Merge Risk: ⚪ Minimal · up to 9ecfd

The token-accounting changes have no identified merge-blocking risk.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 63.64% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the primary billing change: capturing provider usage before response-body parsing. It is concise and specific.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/capture-usage-before-parsing

---

_{Comment @coderabbitai help to get the list of available commands.}