ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: tighten pricing prefix match + clamp corrupt usage

Browse source 
Three follow-ups on the longest-prefix pricing match landed in this PR:

- Prefix match now requires a dash boundary or end-of-string. The
  longest-key sort alone still falsely landed "claude-opus-4-15" on
  the "claude-opus-4-1" row, and "gpt-5.5-prod" on the "gpt-5.5-pro"
  row (a 6x overcharge). Demanding the next character be "-" rules
  out the lookalikes while keeping dated snapshots
  ("gpt-5.4-mini-2026-04-23", "claude-opus-4-7-20260414") landing on
  their canonical row.
- Clamp every token count to >= 0. A corrupted upstream payload
  (negative cached count, off-by-one in a fixture) could previously
  produce a negative bill that masked real spend in the session
  total tooltip.
- Tolerate a non-dict "cache_creation" (e.g. an upstream proxy
  folded the field down to a single int). The current code raised
  AttributeError mid-turn; now it falls back to the 5m-default
  bucket so the rest of the cost calculation still runs.

Adds tests/test_pricing_edge.py with 20 adversarial cases covering
the boundary check, negative / None / zero token values across both
envelopes, cache_read > prompt corruption, the OpenAI long-context
threshold crossover on cache-inflated billable input, malformed
sub-objects, and unknown-provider degradation. Combined suite is
51 tests, all green.
This commit is contained in:
Daniel Han 2026-05-24 14:08:02 +00:00
commit b4bce2d69b
2 changed files with 424 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

39
studio/backend/core/inference/pricing.py
View file

@ -147,8 +147,18 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]:
# like "gpt-5.4-mini-2026-..." match "gpt-5.4-mini" before they
# collide with the shorter "gpt-5.4" entry. Sorting keys by
# length descending picks the most specific table row first.
#
# The match has to land on a token boundary -- otherwise a future
# "claude-opus-4-15" would falsely inherit "claude-opus-4-1" prices,
# and a hypothetical "gpt-5.5-prod" would silently bill as
# "gpt-5.5-pro" (5x overcharge). We require the next character after
# the prefix to be "-" so we only match on dash-separated
# date / variant suffixes, never on a longer numeric or alpha
# continuation of the same component.
for key in sorted(table, key = len, reverse = True):
if model.startswith(key):
if model.startswith(key) and (
len(model) == len(key) or model[len(key)] == "-"
):
return table[key]
return None
@ -206,17 +216,21 @@ def calculate_cost(
# raw OpenAI: input_tokens INCLUDES cache_read (no cache_create)
# Studio Anthropic: prompt_tokens INCLUDES cache_creation + cache_read
# Studio OpenAI: prompt_tokens == raw input_tokens (includes cache_read)
cache_creation = int(usage.get("cache_creation_input_tokens") or 0)
cache_read = int(usage.get("cache_read_input_tokens") or 0)
# Defense in depth: clamp every token count to >=0 so a corrupted
# upstream payload (negative cached count, off-by-one in a fixture)
# can never produce a NEGATIVE bill that masks real spend in the
# session total tooltip.
cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0))
cache_read = max(0, int(usage.get("cache_read_input_tokens") or 0))
has_input_tokens = "input_tokens" in usage and usage.get("input_tokens") is not None
if has_input_tokens:
# Raw upstream envelope.
input_tokens = int(usage.get("input_tokens") or 0)
input_tokens = max(0, int(usage.get("input_tokens") or 0))
else:
# Studio chat-style envelope: prompt_tokens already folds the
# cache buckets for Anthropic, so peel them off to recover the
# raw uncached prompt count and keep downstream math symmetric.
prompt_tokens = int(usage.get("prompt_tokens") or 0)
prompt_tokens = max(0, int(usage.get("prompt_tokens") or 0))
if provider == "anthropic":
input_tokens = max(0, prompt_tokens - cache_creation - cache_read)
else:
@ -226,9 +240,9 @@ def calculate_cost(
# ``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)
output_tokens = max(0, int(usage.get("output_tokens") or 0))
else:
output_tokens = int(usage.get("completion_tokens") or 0)
output_tokens = max(0, int(usage.get("completion_tokens") or 0))
if provider == "openai":
# Cached prompt tokens live on different sub-objects depending on
# which envelope landed here. Raw OpenAI Responses usage uses
@ -277,10 +291,13 @@ def calculate_cost(
if provider == "anthropic":
# Split cache_creation across 5m / 1h buckets when the
# response surfaces the breakdown.
cc_breakdown = usage.get("cache_creation") or {}
cc_5m = int(cc_breakdown.get("ephemeral_5m_input_tokens") or 0)
cc_1h = int(cc_breakdown.get("ephemeral_1h_input_tokens") or 0)
# response surfaces the breakdown. Tolerate a non-dict value
# (some upstream proxies fold the field down to an int total)
# so we don't crash the cost calculation on a malformed payload.
cc_raw = usage.get("cache_creation")
cc_breakdown = cc_raw if isinstance(cc_raw, dict) else {}
cc_5m = max(0, int(cc_breakdown.get("ephemeral_5m_input_tokens") or 0))
cc_1h = max(0, int(cc_breakdown.get("ephemeral_1h_input_tokens") or 0))
if cc_5m + cc_1h == 0 and cache_creation > 0:
# Fall back: assume default 5m pool when no breakdown is given.
cc_5m = cache_creation

396
studio/backend/tests/test_pricing_edge.py Normal file
View file

@ -0,0 +1,396 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Adversarial edge cases for ``calculate_cost`` / ``_lookup``.
The companion ``test_pricing.py`` module pins the happy-path math.
This module pins the corner cases reviewers asked about on the
pricing follow-up:
- prefix matching must land on a token boundary so a future
``claude-opus-4-15`` can't silently inherit ``claude-opus-4-1``
pricing, and a hypothetical ``gpt-5.5-prod`` can't bill as the
much more expensive ``gpt-5.5-pro``;
- negative / non-int token values from a corrupted upstream payload
must never produce a negative bill that masks real spend in the
running session total;
- chat-style and raw envelopes must produce identical totals on the
cache-heavy boundary cases (zero cache_read, cache_read > prompt,
cache_creation > prompt) so the cost ledger doesn't flip per
envelope choice;
- the long-context tier crossover has to fire on the ``billable``
input count (which folds cache_creation for OpenAI), not the raw
``input_tokens`` only, otherwise cache-heavy turns dodge the
long-context premium they should pay;
- malformed sub-objects (cache_creation as int, server_tool_use as
str) must not raise -- the calculator surfaces zero cost for the
malformed bucket and keeps pricing the rest of the turn.
"""
import math
from core.inference.pricing import (
ANTHROPIC_CACHE_5M_WRITE_MULT,
ANTHROPIC_CACHE_READ_MULT,
ANTHROPIC_PRICING,
OPENAI_CACHE_READ_MULT,
OPENAI_PRICING,
_lookup,
calculate_cost,
)
def _isclose(a, b, tol = 1e-6):
return math.isclose(a, b, rel_tol = tol, abs_tol = tol)
# ── prefix-match boundary checks ────────────────────────────────────
def test_prefix_match_requires_dash_boundary_opus_variant():
# Hypothetical future ``claude-opus-4-15`` must NOT inherit the
# ``claude-opus-4-1`` ($15 / $75) row by way of a naive substring
# ``startswith`` -- the next char has to be ``-`` or end-of-string.
assert _lookup("anthropic", "claude-opus-4-15") is None
out = calculate_cost(
"anthropic",
"claude-opus-4-15",
{"input_tokens": 1_000_000, "output_tokens": 0},
)
assert out["priced"] is False
assert out["total_usd"] == 0.0
def test_prefix_match_requires_dash_boundary_gpt_variant():
# ``gpt-5.55`` is hypothetical but defends the same invariant:
# a different model in the same family tree can't piggy-back on the
# ``gpt-5.5`` row just because the canonical id is a string prefix.
assert _lookup("openai", "gpt-5.55") is None
assert _lookup("openai", "gpt-5.55-2026-04-23") is None
out = calculate_cost(
"openai",
"gpt-5.55-2026-04-23",
{"input_tokens": 1_000_000, "output_tokens": 0},
)
assert out["priced"] is False
def test_prefix_match_requires_dash_boundary_pro_lookalike():
# ``gpt-5.5-pro`` and ``gpt-5.5-prod`` share an undelimited "prod"
# / "pro" component. A naive prefix match would land
# ``gpt-5.5-prod`` on the ``gpt-5.5-pro`` row ($30 / MTok) -- a 6x
# overcharge against the canonical ``gpt-5.5`` row ($5 / MTok). The
# dash-boundary check forces the longest matching key to end on
# ``-``, so ``gpt-5.5-prod`` falls through ``gpt-5.5-pro`` and
# lands on the canonical ``gpt-5.5`` row instead.
prices = _lookup("openai", "gpt-5.5-prod")
assert prices is not None
assert prices["input_per_mtok"] == OPENAI_PRICING["gpt-5.5"]["input_per_mtok"], (
"expected fallback to gpt-5.5 base ($5), not gpt-5.5-pro ($30)"
)
out = calculate_cost(
"openai",
"gpt-5.5-prod",
{"input_tokens": 100_000, "output_tokens": 0},
)
assert out["priced"] is True
assert _isclose(out["input_usd"], 100_000 / 1_000_000.0 * 5.0)
def test_prefix_match_still_resolves_legit_dated_snapshots():
# The boundary fix must NOT regress the real win the PR is for.
# gpt-5.4-mini-2026-04-23 inherits the mini row, NOT gpt-5.4's.
out = calculate_cost(
"openai",
"gpt-5.4-mini-2026-04-23",
{"input_tokens": 1_000_000, "output_tokens": 0},
)
assert out["priced"] is True
assert _isclose(out["input_usd"], 0.75)
# And an Anthropic dated snapshot still resolves to its canonical row.
out = calculate_cost(
"anthropic",
"claude-opus-4-7-20260414",
{"input_tokens": 1_000_000, "output_tokens": 0},
)
assert out["priced"] is True
assert _isclose(out["input_usd"], 5.0)
# ── precedence: input_tokens wins over prompt_tokens (and 0 is real) ──
def test_explicit_zero_input_tokens_wins_over_stale_prompt_tokens():
# Mirror of ``test_explicit_zero_output_tokens_wins_over_stale_completion_tokens``
# for the input side -- a turn that genuinely consumed zero prompt
# tokens must not silently bill against the chat-style mirror.
out = calculate_cost(
"openai",
"gpt-5.5",
{
"input_tokens": 0,
"prompt_tokens": 1_000_000, # stale chat-style mirror
"output_tokens": 100,
},
)
assert out["billable_input_tokens"] == 0
assert out["input_usd"] == 0.0
def test_none_input_tokens_falls_through_to_prompt_tokens():
# ``None`` is the "key present but unset" case -- treated as
# missing so the chat-style mirror still wins.
out = calculate_cost(
"openai",
"gpt-5.5",
{
"input_tokens": None,
"prompt_tokens": 200_000,
"output_tokens": None,
"completion_tokens": 5_000,
},
)
assert out["billable_input_tokens"] == 200_000
assert out["billable_output_tokens"] == 5_000
assert _isclose(out["input_usd"], 200_000 / 1_000_000.0 * 5.0)
assert _isclose(out["output_usd"], 5_000 / 1_000_000.0 * 30.0)
# ── negative / corrupted upstream values clamp to zero ──────────────
def test_negative_tokens_clamp_to_zero_no_negative_bill():
out = calculate_cost(
"openai",
"gpt-5.5",
{"input_tokens": -100, "output_tokens": -50},
)
assert out["billable_input_tokens"] == 0
assert out["billable_output_tokens"] == 0
assert out["input_usd"] == 0.0
assert out["output_usd"] == 0.0
assert out["total_usd"] == 0.0
def test_negative_cache_buckets_clamp_to_zero():
# Negative cache_read on Anthropic would otherwise refund the bill.
out = calculate_cost(
"anthropic",
"claude-opus-4-7",
{
"input_tokens": 1_000,
"output_tokens": 0,
"cache_creation_input_tokens": -500,
"cache_read_input_tokens": -1_000,
},
)
assert out["cache_write_usd"] == 0.0
assert out["cache_read_usd"] == 0.0
assert out["billable_input_tokens"] == 1_000
assert out["total_usd"] >= 0.0
def test_negative_prompt_tokens_chat_style_clamp():
out = calculate_cost(
"openai",
"gpt-5.4-mini",
{"prompt_tokens": -100, "completion_tokens": -50},
)
assert out["billable_input_tokens"] == 0
assert out["billable_output_tokens"] == 0
assert out["total_usd"] == 0.0
# ── cache_read > prompt_tokens corruption: no negative billable ─────
def test_anthropic_chat_cache_read_exceeds_prompt_no_negative_billable():
# If cache_read claims more tokens than prompt_tokens reports,
# the uncached_input should clamp at 0 -- but the billable counter
# still reflects the cache buckets (we charge for what we got).
out = calculate_cost(
"anthropic",
"claude-opus-4-7",
{
"prompt_tokens": 100,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 500,
"completion_tokens": 0,
},
)
assert out["input_usd"] == 0.0 # uncached clamped to 0
assert out["billable_input_tokens"] == 500 # 0 uncached + 500 cache_read
# cache_read still priced at the discount rate.
base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"]
assert _isclose(
out["cache_read_usd"], 500 / 1_000_000.0 * base * ANTHROPIC_CACHE_READ_MULT
)
def test_openai_raw_cached_tokens_exceeds_input_clamp_non_cached():
# OpenAI variant of the same defense: cached > input shouldn't
# produce a negative "non_cached_input" or a negative input_usd.
base = OPENAI_PRICING["gpt-5.5"]["input_per_mtok"]
out = calculate_cost(
"openai",
"gpt-5.5",
{
"input_tokens": 100,
"output_tokens": 0,
"input_tokens_details": {"cached_tokens": 500},
},
)
assert out["input_usd"] == 0.0
# Cache read still priced (the 0.1x bucket).
assert _isclose(
out["cache_read_usd"], 500 / 1_000_000.0 * base * OPENAI_CACHE_READ_MULT
)
# ── long-context tier crosses on billable, including cache_creation ──
def test_openai_long_context_triggers_on_cache_creation_inflated_billable():
# 250k raw input + 50k cache_creation pushes billable to 300k, which
# crosses the 272k threshold. The whole turn should price at the
# long-context tier so we don't undercount.
out = calculate_cost(
"openai",
"gpt-5.5",
{
"input_tokens": 250_000,
"cache_creation_input_tokens": 50_000,
"output_tokens": 1_000,
},
)
assert out["billable_input_tokens"] == 300_000
assert "long-context" in out["model_priced"]
assert _isclose(out["input_usd"], 250_000 / 1_000_000.0 * 10.0)
assert _isclose(out["output_usd"], 1_000 / 1_000_000.0 * 45.0)
def test_openai_long_context_threshold_boundary_inclusive():
# 272_000 exactly should land in the long-context tier (>=).
out = calculate_cost(
"openai",
"gpt-5.5",
{"input_tokens": 272_000, "output_tokens": 1_000},
)
assert "long-context" in out["model_priced"]
out_lo = calculate_cost(
"openai",
"gpt-5.5",
{"input_tokens": 271_999, "output_tokens": 1_000},
)
assert "long-context" not in out_lo["model_priced"]
# ── chat-style vs raw envelope parity at OpenAI long-context tier ──
def test_openai_chat_envelope_long_context_parity_with_raw():
raw = calculate_cost(
"openai",
"gpt-5.5",
{"input_tokens": 300_000, "output_tokens": 10_000},
)
chat = calculate_cost(
"openai",
"gpt-5.5",
{"prompt_tokens": 300_000, "completion_tokens": 10_000},
)
assert _isclose(chat["total_usd"], raw["total_usd"])
assert "long-context" in chat["model_priced"]
assert "long-context" in raw["model_priced"]
# ── malformed sub-objects: no crash, no false bill ──────────────────
def test_cache_creation_as_int_does_not_crash():
# Some upstream proxies fold cache_creation down to a single int.
# The calculator must tolerate that and fall back to the 5m default.
base = ANTHROPIC_PRICING["claude-opus-4-7"]["input_per_mtok"]
out = calculate_cost(
"anthropic",
"claude-opus-4-7",
{
"input_tokens": 0,
"output_tokens": 0,
"cache_creation_input_tokens": 1_000_000,
"cache_creation": 12345, # malformed; must not raise
},
)
# Falls back to the 5m default for the whole cache_creation bucket.
assert _isclose(
out["cache_write_usd"], 1_000_000 / 1_000_000.0 * base * ANTHROPIC_CACHE_5M_WRITE_MULT
)
def test_non_dict_server_tool_use_is_ignored():
out = calculate_cost(
"anthropic",
"claude-opus-4-7",
{"input_tokens": 100, "output_tokens": 100, "server_tool_use": "garbage"},
)
assert out["server_tools_usd"] == 0.0
out = calculate_cost(
"openai",
"gpt-5.5",
{"input_tokens": 100, "output_tokens": 100, "openai_tool_use": [1, 2, 3]},
)
assert out["server_tools_usd"] == 0.0
def test_non_dict_input_tokens_details_is_ignored():
out = calculate_cost(
"openai",
"gpt-5.5",
{
"input_tokens": 100,
"output_tokens": 0,
"input_tokens_details": "nope",
"prompt_tokens_details": [1, 2, 3],
},
)
# No cached_tokens recovered, so no cache_read_usd discount.
assert out["cache_read_usd"] == 0.0
# ── unknown provider degrades gracefully ────────────────────────────
def test_unknown_provider_priced_false_zero_bill():
out = calculate_cost(
"gemini",
"gemini-pro",
{"input_tokens": 1_000_000, "output_tokens": 1_000_000},
)
assert out["priced"] is False
assert out["total_usd"] == 0.0
# Tokens still report for the UI counter.
assert out["billable_input_tokens"] == 1_000_000
assert out["billable_output_tokens"] == 1_000_000
def test_anthropic_provider_with_openai_model_priced_false():
# Provider routing mistake: handing an OpenAI id to the Anthropic
# table must not falsely match a similar-looking Anthropic key.
out = calculate_cost(
"anthropic",
"gpt-5.5",
{"input_tokens": 1_000_000, "output_tokens": 0},
)
assert out["priced"] is False
# ── all-zero / empty usage stays at zero ────────────────────────────
def test_empty_usage_dict_zero_bill():
out = calculate_cost("openai", "gpt-5.5", {})
assert out["priced"] is True # model is in the table
assert out["billable_input_tokens"] == 0
assert out["total_usd"] == 0.0