From 0bd831dc5e17f0053e2639fde1cd29dd08cb982d Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Sat, 11 Jul 2026 14:55:32 -0700 Subject: [PATCH] fix(proxy): reserve tiered output at the higher reasoning rate Some tiered Dashscope models price reasoning output above standard output (output_cost_per_reasoning_token > output_cost_per_token). The reservation charged all output at the standard rate, so a reasoning-heavy request reserved too little and concurrent calls could exceed the budget before reconciliation. The reasoning share is unknown before the request runs, so reserve every output token at the higher of the two configured rates, for both tiered and flat pricing. Co-authored-by: Cursor --- .../spend_tracking/budget_reservation.py | 15 +++- .../proxy/test_budget_reservation.py | 81 +++++++++++++++++++ 2 files changed, 93 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f02b733229a..38fd0d3f343 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1014,20 +1014,29 @@ def _max_cost_for_cost_info( if isinstance(tiered_pricing, list) and tiered_pricing: tier = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=input_tokens) if tier is not None: + output_rate = max( + tier_rate(tier, "output_cost_per_token"), + tier_rate(tier, "output_cost_per_reasoning_token"), + ) return (input_tokens * tier_rate(tier, "input_cost_per_token")) + ( - output_tokens * output_multiplier * tier_rate(tier, "output_cost_per_token") + output_tokens * output_multiplier * output_rate ) input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) output_cost_per_token = _to_float(model_info.get("output_cost_per_token")) + output_cost_per_reasoning_token = _to_float(model_info.get("output_cost_per_reasoning_token")) cost = 0.0 if input_cost_per_token is not None: cost += input_tokens * input_cost_per_token elif input_tokens > 0: return None - if output_cost_per_token is not None: - cost += output_tokens * output_multiplier * output_cost_per_token + # The reasoning-token share is unknown before the request runs, so reserve every + # output token at the higher of the standard and reasoning rates to avoid + # under-reserving reasoning-heavy requests. + output_rate = max(output_cost_per_token or 0.0, output_cost_per_reasoning_token or 0.0) + if output_cost_per_token is not None or output_cost_per_reasoning_token is not None: + cost += output_tokens * output_multiplier * output_rate elif output_tokens > 0: return None diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index efc4c9758bf..3da103683ba 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -868,6 +868,87 @@ def test_tiered_reservation_is_all_or_nothing_with_output_tier_from_input_length assert estimated > graduated_under_reserve +def test_tiered_reservation_uses_higher_reasoning_output_rate(): + """Some tiered models price reasoning output above standard output. The + reasoning-token share is unknown before the request runs, so reservation must + charge every output token at the higher of the two rates to avoid under-reserving + reasoning-heavy requests.""" + tiered_pricing = [ + { + "range": [0, 32000], + "input_cost_per_token": 1e-06, + "output_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 4e-06, + } + ] + input_tokens = 1000 + output_tokens = 500 + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={"tiered_pricing": tiered_pricing, "max_output_tokens": 200000}, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_input_tokens", + return_value=input_tokens, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_output_tokens", + return_value=output_tokens, + ), + ): + estimated = estimate_request_max_cost( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + ) + + expected = (input_tokens * 1e-06) + (output_tokens * 4e-06) + assert estimated == pytest.approx(expected) + + # Reserving output at the plain rate would under-reserve reasoning-heavy calls. + under_reserve = (input_tokens * 1e-06) + (output_tokens * 1.2e-06) + assert estimated > under_reserve + + +def test_flat_reservation_uses_higher_reasoning_output_rate(): + """The same reasoning under-reservation gap exists for flat-rate models that + declare output_cost_per_reasoning_token above output_cost_per_token.""" + input_tokens = 1000 + output_tokens = 500 + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "input_cost_per_token": 1e-06, + "output_cost_per_token": 1.2e-06, + "output_cost_per_reasoning_token": 4e-06, + "max_output_tokens": 200000, + }, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_input_tokens", + return_value=input_tokens, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation._estimate_output_tokens", + return_value=output_tokens, + ), + ): + estimated = estimate_request_max_cost( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + ) + + expected = (input_tokens * 1e-06) + (output_tokens * 4e-06) + assert estimated == pytest.approx(expected) + under_reserve = (input_tokens * 1e-06) + (output_tokens * 1.2e-06) + assert estimated > under_reserve + + def test_reservation_uses_most_expensive_deployment_in_group(): """When a model group mixes deployments with different tiered rates, reservation must estimate against the most expensive one. Reserving the cheaper sibling would