fix: guard cache_creation_tokens overwrite + add dict-path test

- Guard against silently overwriting a pre-existing cache_creation_tokens
  value on the PromptTokensDetailsWrapper when the OpenRouter block fires.
  If cache_creation_tokens is already set (e.g. by a prior Anthropic-native
  mapping pass), the cache_write_tokens value must not replace it.

- Add test_usage_openrouter_cache_tokens_dict_path: covers the plain-dict
  branch (isinstance(prompt_tokens_details, dict) →
  PromptTokensDetailsWrapper(**d)), which is the actual production code path
  when Usage is constructed from a JSON-parsed API response. Asserts both
  private fields and cache_creation_tokens are correctly populated.

- Extend existing test with a no-overwrite assertion.

Made-with: Cursor
This commit is contained in:
drexkooo 2026-03-17 14:21:47 +01:00
parent db23e717c7
commit aaf492af17
2 changed files with 44 additions and 2 deletions

View file

@ -1681,8 +1681,11 @@ class Usage(SafeAttributeModel, CompletionUsage):
self._cache_creation_input_tokens = _writes
# Also populate the public field that the cost calculator reads
# (_parse_prompt_tokens_details reads cache_creation_tokens, not
# the private _cache_creation_input_tokens attribute)
_ptd.cache_creation_tokens = _writes
# the private _cache_creation_input_tokens attribute).
# Guard against overwriting a value already set by a prior
# Anthropic-native mapping pass.
if not getattr(_ptd, "cache_creation_tokens", None):
_ptd.cache_creation_tokens = _writes
for k, v in params.items():
setattr(self, k, v)

View file

@ -361,3 +361,42 @@ def test_usage_openrouter_cache_tokens_from_prompt_tokens_details():
assert "cache_write_tokens" not in ptd
# cache_creation_tokens should also be absent when no writes were reported
assert "cache_creation_tokens" not in ptd
# If cache_creation_tokens is already set on the wrapper (e.g. from a prior
# Anthropic-native mapping pass), the OpenRouter block must not overwrite it.
usage_no_overwrite = Usage(
prompt_tokens=18500,
completion_tokens=120,
total_tokens=18620,
prompt_tokens_details=PromptTokensDetailsWrapper(
cache_creation_tokens=999, # pre-existing value
cache_write_tokens=400, # OpenRouter field — must NOT win
),
)
assert usage_no_overwrite.prompt_tokens_details.cache_creation_tokens == 999
def test_usage_openrouter_cache_tokens_dict_path():
"""In production, prompt_tokens_details arrives as a plain dict (parsed from JSON).
The dict branch (isinstance(..., dict) PromptTokensDetailsWrapper(**d)) must also
populate both the private fields and cache_creation_tokens for cost calculation."""
from litellm.types.utils import Usage
# Plain dict — the real code path from JSON-parsed API responses
usage = Usage(
prompt_tokens=18500,
completion_tokens=120,
total_tokens=18620,
prompt_tokens_details={
"cached_tokens": 17000,
"cache_write_tokens": 400,
},
)
# Both the streaming-adapter path (private fields) …
assert usage._cache_read_input_tokens == 17000
assert usage._cache_creation_input_tokens == 400
# … and the cost-calculator path (public wrapper field) must be populated.
assert usage.prompt_tokens_details is not None
assert usage.prompt_tokens_details.cache_creation_tokens == 400