From 8ac670d77a4e005222b56304a170e88a164fd782 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 22 May 2026 17:13:22 +0000
Subject: [PATCH] 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.
---
 studio/backend/core/inference/pricing.py | 11 ++++++++---
 studio/backend/tests/test_pricing.py     | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/pricing.py b/studio/backend/core/inference/pricing.py
index dc021dc020..d9033e73ea 100644
--- a/studio/backend/core/inference/pricing.py
+++ b/studio/backend/core/inference/pricing.py
@@ -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):
diff --git a/studio/backend/tests/test_pricing.py b/studio/backend/tests/test_pricing.py
index 704e4a4fc7..662866187f 100644
--- a/studio/backend/tests/test_pricing.py
+++ b/studio/backend/tests/test_pricing.py
@@ -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