Studio: prefer raw output_tokens over chat-style completion_tokens

Codex flagged that the previous fallback chain
'usage.get("output_tokens") or usage.get("completion_tokens")'
treats an explicit 0 as missing -- a mixed-envelope payload where
'output_tokens' is 0 but 'completion_tokens' is non-zero (or
stale) bills the wrong amount. Mirror the has_input_tokens
precedence pattern: when the raw key is present we use it even at
0; otherwise fall back to completion_tokens.
This commit is contained in:
Daniel Han 2026-05-22 17:13:22 +00:00 committed by danielhanchen
commit 8ac670d77a
2 changed files with 30 additions and 3 deletions

View file

@ -221,9 +221,14 @@ def calculate_cost(
input_tokens = max(0, prompt_tokens - cache_creation - cache_read)
else:
input_tokens = prompt_tokens
output_tokens = int(
usage.get("output_tokens") or usage.get("completion_tokens") or 0
)
# Prefer the raw upstream key when present, even when its value is
# explicitly 0 -- the ``or`` fallback would mistakenly pick a stale
# ``completion_tokens`` for an empty completion. Mirrors the
# has_input_tokens precedence above.
if "output_tokens" in usage and usage.get("output_tokens") is not None:
output_tokens = int(usage.get("output_tokens") or 0)
else:
output_tokens = int(usage.get("completion_tokens") or 0)
if provider == "openai":
details = usage.get("input_tokens_details") or {}
if isinstance(details, dict):

View file

@ -552,3 +552,25 @@ def test_openai_chat_style_prompt_tokens_keeps_cache_read_semantics():
},
)
assert _isclose(chat["total_usd"], raw["total_usd"]), (chat, raw)
def test_explicit_zero_output_tokens_wins_over_stale_completion_tokens():
"""``output_tokens: 0`` must beat a stale ``completion_tokens: 50``.
The previous ``or`` fallback treated 0 as missing and silently
re-priced the response against the stale chat-style count. The
raw upstream key now takes precedence even when its value is 0.
"""
out = calculate_cost(
"openai",
"gpt-4o-mini",
{
"input_tokens": 100,
"output_tokens": 0,
# Mixed envelope: stale chat-style completion_tokens that
# the caller forgot to clear. We must NOT bill against it.
"completion_tokens": 50,
},
)
assert out["billable_output_tokens"] == 0, out
assert out["output_usd"] == 0.0, out