From adc41ade8c0eb15ec6165906b2c3c9803b26d386 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 8 May 2026 20:18:31 -0700 Subject: [PATCH 1/3] fix(proxy): bound budget reservation per request instead of pinning to remaining headroom reserve_budget_for_request fell back to reserving the entire remaining team/key/user headroom whenever a request omitted max_tokens, which pinned the spend counter at max_budget for the duration of the in-flight request and false-positive-blocked every concurrent or back-to-back request until the success callback reconciled. Surfaced as an integration-test team being budget-blocked at its $2000 cap while DB spend was $0.144. Switch the missing-max_tokens path to a fixed default of 16384 output tokens (mirrors parallel_request_limiter_v3's DEFAULT_MAX_TOKENS_ESTIMATE precedent), and clamp explicit max_tokens at the model's max_output_tokens for reservation accounting only. The outbound request body is unchanged, so providers see whatever the caller actually sent; only the local integer used to compute reservation cost is bounded. This also prevents a hostile max_tokens=999999999 from inflating one request's reservation up to the entire team headroom. For Opus 4.7 (output $25/M, max_output 128K) on a $2000 budget the worst-case per-request reservation drops from "everything left" to $3.20, raising admittable concurrency from 1 to ~625. --- .../spend_tracking/budget_reservation.py | 63 ++-- .../proxy/test_budget_reservation.py | 327 ++++++++---------- 2 files changed, 164 insertions(+), 226 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 1d296611bf..feae368d23 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -95,11 +95,9 @@ async def reserve_budget_for_request( route=route, llm_router=llm_router, ) - if reservation_cost is None: - reservation_cost = await _get_smallest_remaining_budget( - counters=counters, - current_spend_by_counter_key=current_spend_by_counter_key, - ) + # estimate_request_max_cost still returns None when the model is unknown + # to the cost map (no token-priced cost fields, e.g. image/audio routes). + # In that case we fall back to read-time enforcement only. if reservation_cost is None or reservation_cost <= 0: return None @@ -553,32 +551,6 @@ def _coerce_window(window: Any) -> dict: return {} -async def _get_smallest_remaining_budget( - counters: List[_BudgetCounter], - current_spend_by_counter_key: Dict[str, float], -) -> Optional[float]: - remaining_budget: Optional[float] = None - for counter in counters: - current_spend = await _get_current_counter_value(counter=counter) - current_spend_by_counter_key[counter.counter_key] = current_spend - remaining = counter.max_budget - current_spend - if remaining <= 0: - raise litellm.BudgetExceededError( - current_cost=current_spend, - max_budget=counter.max_budget, - message=( - "Budget has been exceeded! " - f"{counter.entity_type}={counter.entity_id} " - f"Current cost: {current_spend}, " - f"Max budget: {counter.max_budget}" - ), - ) - remaining_budget = ( - remaining if remaining_budget is None else min(remaining_budget, remaining) - ) - return remaining_budget - - async def _reserve_counter( counter: _BudgetCounter, reservation_cost: float, @@ -946,6 +918,9 @@ def _estimate_input_tokens( return None +DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK = 16384 + + def _estimate_output_tokens( request_body: dict, route: str, @@ -954,15 +929,27 @@ def _estimate_output_tokens( if _is_input_only_route(route=route): return 0 + requested: Optional[int] = None for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): - max_tokens = _to_int(request_body.get(key)) - if max_tokens is not None: - return max_tokens + requested = _to_int(request_body.get(key)) + if requested is not None: + break - # If the caller did not cap output tokens, avoid reserving a model's - # theoretical maximum context. The caller can still admit one request by - # reserving the smallest remaining budget in reserve_budget_for_request(). - return None + # Clamp at min(requested-or-default, model_max-or-default). Two purposes: + # (1) Without an explicit cap we still need a finite reservation so the + # atomic admission counter actually bounds concurrent in-flight cost + # (mirrors parallel_request_limiter_v3's DEFAULT_MAX_TOKENS_ESTIMATE). + # (2) An adversarial caller cannot send max_tokens=999999999 to inflate + # the reservation up to remaining team headroom and pin the counter + # at the cap — the model can only physically emit max_output_tokens + # anyway, so reserving more is both wasteful and a DoS surface. + model_ceiling = ( + _to_int(model_info.get("max_output_tokens")) + or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK + ) + if requested is None: + requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK + return min(requested, model_ceiling) def _count_text_tokens(model: str, text: Any) -> int: diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 070b232066..0e2ca98a11 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -585,15 +585,24 @@ async def test_should_cap_known_estimate_to_remaining_budget( @pytest.mark.asyncio -async def test_should_reserve_remaining_budget_when_output_cap_missing( +async def test_should_clamp_reservation_to_default_when_output_cap_missing( spend_counter_state, ): + """When max_tokens is not specified, _estimate_output_tokens falls back to + DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK (16K), clamped by the model's + max_output_tokens. Reservation must be a bounded per-request amount + (mirroring parallel_request_limiter_v3's DEFAULT_MAX_TOKENS_ESTIMATE), + not the entire remaining headroom.""" + from litellm.proxy.spend_tracking.budget_reservation import ( + DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK, + ) + counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) valid_token = UserAPIKeyAuth( token="key-budget-uncapped", spend=0.2, - max_budget=1.0, + max_budget=10000.0, ) await key_cache.async_set_cache( key="key-budget-uncapped", @@ -602,22 +611,24 @@ async def test_should_reserve_remaining_budget_when_output_cap_missing( request_body = _request_body() request_body.pop("max_tokens") + output_cost_per_token = 1e-5 # roughly Opus 4.5/4.7 output rate + expected_cost = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK * output_cost_per_token + with patch( "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", return_value={ "input_cost_per_token": 0.0, - "output_cost_per_token": 100.0, - "max_output_tokens": 200000, + "output_cost_per_token": output_cost_per_token, + "max_output_tokens": 200000, # well above the 16K fallback }, ): - assert ( - estimate_request_max_cost( - request_body=request_body, - route="/chat/completions", - llm_router=None, - ) - is None + estimated = estimate_request_max_cost( + request_body=request_body, + route="/chat/completions", + llm_router=None, ) + assert estimated == pytest.approx(expected_cost) + reservation = await reserve_budget_for_request( request_body=request_body, route="/chat/completions", @@ -631,47 +642,45 @@ async def test_should_reserve_remaining_budget_when_output_cap_missing( ) assert reservation is not None - assert reservation["reserved_cost"] == pytest.approx(0.8) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-uncapped" - ) == pytest.approx(1.0) - + assert reservation["reserved_cost"] == pytest.approx(expected_cost) await release_budget_reservation(reservation) @pytest.mark.asyncio -async def test_should_shrink_uncapped_reservation_when_counter_advances( +async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequests( spend_counter_state, - monkeypatch, ): + """An adversarial caller sending max_tokens=999_999_999 must not be able + to inflate the per-request reservation up to the entire remaining team + headroom. _estimate_output_tokens clamps the explicit value at the + model's max_output_tokens — the model can only physically emit that + many tokens anyway, so anything more is both wasteful and a DoS surface.""" counter_cache, key_cache = spend_counter_state proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) valid_token = UserAPIKeyAuth( - token="key-budget-uncapped-race", - spend=0.2, - max_budget=1.0, + token="key-budget-overrequest", + spend=0.0, + max_budget=10000.0, ) + await key_cache.async_set_cache( + key="key-budget-overrequest", + value=valid_token, + ) + request_body = _request_body() - request_body.pop("max_tokens") + request_body["max_tokens"] = 999_999_999 - from litellm.proxy.spend_tracking import budget_reservation - - async def stale_counter_read(counter): - await counter_cache.async_increment_cache( - key=counter.counter_key, - value=0.3, - ) - return 0.2 - - monkeypatch.setattr( - budget_reservation, - "_get_current_counter_value", - stale_counter_read, - ) + output_cost_per_token = 1e-5 + model_ceiling = 128_000 + expected_cost = model_ceiling * output_cost_per_token with patch( - "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", - return_value=None, + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "input_cost_per_token": 0.0, + "output_cost_per_token": output_cost_per_token, + "max_output_tokens": model_ceiling, + }, ): reservation = await reserve_budget_for_request( request_body=request_body, @@ -686,101 +695,9 @@ async def test_should_shrink_uncapped_reservation_when_counter_advances( ) assert reservation is not None - assert reservation["reserved_cost"] == pytest.approx(0.7) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-uncapped-race" - ) == pytest.approx(1.0) - + assert reservation["reserved_cost"] == pytest.approx(expected_cost) await release_budget_reservation(reservation) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-uncapped-race" - ) == pytest.approx(0.3) - - -@pytest.mark.asyncio -async def test_should_shrink_uncapped_reservation_multiple_times( - spend_counter_state, - monkeypatch, -): - counter_cache, key_cache = spend_counter_state - proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) - valid_token = UserAPIKeyAuth( - token="key-budget-double-resize", - spend=0.2, - max_budget=1.0, - team_id="team-budget-double-resize", - ) - team_object = LiteLLM_TeamTable( - team_id="team-budget-double-resize", - spend=0.2, - max_budget=1.0, - ) - request_body = _request_body() - request_body.pop("max_tokens") - - from litellm.proxy.spend_tracking import budget_reservation - - stale_spend_by_counter_key = { - "spend:key:key-budget-double-resize": 0.3, - "spend:team:team-budget-double-resize": 0.4, - } - - async def stale_counter_read(counter): - await counter_cache.async_increment_cache( - key=counter.counter_key, - value=stale_spend_by_counter_key[counter.counter_key], - ) - return 0.2 - - monkeypatch.setattr( - budget_reservation, - "_get_current_counter_value", - stale_counter_read, - ) - - with patch( - "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", - return_value=None, - ): - reservation = await reserve_budget_for_request( - request_body=request_body, - route="/chat/completions", - llm_router=None, - valid_token=valid_token, - team_object=team_object, - user_object=None, - prisma_client=None, - user_api_key_cache=key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - assert reservation is not None - assert reservation["reserved_cost"] == pytest.approx(0.6) - assert [entry["reserved_cost"] for entry in reservation["entries"]] == [ - pytest.approx(0.6), - pytest.approx(0.6), - ] - assert [entry["applied_adjustment"] for entry in reservation["entries"]] == [ - pytest.approx(0.0), - pytest.approx(0.0), - ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-double-resize" - ) == pytest.approx(0.9) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team:team-budget-double-resize" - ) == pytest.approx(1.0) - - await release_budget_reservation(reservation) - - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-budget-double-resize" - ) == pytest.approx(0.3) - assert counter_cache.in_memory_cache.get_cache( - key="spend:team:team-budget-double-resize" - ) == pytest.approx(0.4) - def test_should_start_window_without_reset_at_at_duration_boundary(): before = datetime.now(timezone.utc) - timedelta(hours=1) @@ -1047,62 +964,6 @@ async def test_should_release_tracked_entry_when_reservation_fails_after_increme ) == pytest.approx(0.0) -@pytest.mark.asyncio -async def test_should_not_re_read_uncapped_budget_after_reservation_fallback( - spend_counter_state, - monkeypatch, -): - _, key_cache = spend_counter_state - proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) - valid_token = UserAPIKeyAuth( - token="key-budget-uncapped-read-once", - spend=0.2, - max_budget=1.0, - ) - - from litellm.proxy.spend_tracking import budget_reservation - - current_counter_reads = [] - - async def mock_get_current_counter_value(counter): - current_counter_reads.append(counter.counter_key) - return counter.fallback_spend - - async def mock_reserve_counter(counter, reservation_cost): - return None - - monkeypatch.setattr( - budget_reservation, - "_get_current_counter_value", - mock_get_current_counter_value, - ) - monkeypatch.setattr( - budget_reservation, - "_reserve_counter", - mock_reserve_counter, - ) - - with patch( - "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", - return_value=None, - ): - reservation = await reserve_budget_for_request( - request_body=_request_body(), - route="/chat/completions", - llm_router=None, - valid_token=valid_token, - team_object=None, - user_object=None, - prisma_client=None, - user_api_key_cache=key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - assert reservation is not None - assert reservation["reserved_cost"] == pytest.approx(0.8) - assert current_counter_reads == ["spend:key:key-budget-uncapped-read-once"] - - @pytest.mark.asyncio async def test_should_reconcile_reserved_counter_to_actual_spend( spend_counter_state, @@ -1492,4 +1353,94 @@ async def test_should_reserve_all_budgeted_counters(spend_counter_state): counter_cache.in_memory_cache.get_cache(key="spend:team:team-budget-all") == 0.3 ) - await release_budget_reservation(reservation) + +@pytest.mark.asyncio +async def test_should_not_block_concurrent_team_request_when_first_request_lacks_max_tokens( + spend_counter_state, +): + """ + Regression test: a team-bound request with no max_tokens must not pin the + team's spend counter at max_budget for the duration of the request. + + Repro of the integration-test team being falsely budget-blocked at the + $2000 cap while DB spend is $0.144: the first request without max_tokens + used to reserve the entire remaining headroom, leaving any subsequent + request stuck behind a counter sitting at the cap until the success + callback finished reconciling. + """ + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + + valid_token = UserAPIKeyAuth( + token="key-team-integration-tests", + spend=0.0, + team_id="team-integration-tests", + ) + team_object = LiteLLM_TeamTable( + team_id="team-integration-tests", + max_budget=2000.0, + spend=0.144, + ) + await key_cache.async_set_cache( + key=f"team_id:{team_object.team_id}", + value=team_object, + ) + + request_body = _request_body() + request_body.pop("max_tokens") + + # Realistic Opus 4.7 output pricing — the 16K fallback × $25/M ≈ $0.40 + # reservation per request, leaving ~5000 admittable concurrent requests + # against a $2000 team budget. + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "input_cost_per_token": 5e-6, + "output_cost_per_token": 2.5e-5, + "max_output_tokens": 128000, + }, + ): + first_reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # The team counter must not be pinned at max_budget while the first + # request is in flight, otherwise concurrent requests false-positive. + team_counter_after_first = ( + counter_cache.in_memory_cache.get_cache( + key=f"spend:team:{team_object.team_id}" + ) + or 0.0 + ) + assert team_counter_after_first < team_object.max_budget, ( + f"Team counter sat at {team_counter_after_first} after one uncapped " + f"reservation against a {team_object.max_budget} budget — concurrent " + "requests will be falsely blocked." + ) + + # Second request — same shape — must succeed without raising. + second_reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert second_reservation is not None + + if first_reservation is not None: + await release_budget_reservation(first_reservation) + if second_reservation is not None: + await release_budget_reservation(second_reservation) From 0d551ac4f06b62d270694ee877fdeccf329f5528 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 8 May 2026 21:08:55 -0700 Subject: [PATCH 2/3] fix(proxy): reserve per-image cost for image-generation requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Image-generation routes (dall-e-3, flux, etc.) have no per-token output cost so they fell through to the no-reservation read-time-only path. Concurrent image requests against a depleted budget could all pass common_checks (counter exactly at max_budget passes the strict-`>` gate) and reach the provider before reconciliation caught up. Add per-image reservation in _estimate_request_max_cost_for_model: when the model has a per-image cost field, reserve `n × cost_per_image` upfront. The atomic counter increment serializes concurrent admissions, so the second request sees the post-first-reservation counter and raises BudgetExceededError instead of silently leaking through. Both `output_cost_per_image` and `input_cost_per_image` are honored — naming is inconsistent across providers (OpenAI dall-e-3 uses input_cost_per_image, aiml/dall-e-3 uses output_cost_per_image for the same per-generated-image price). Per-pixel pricing (DALL-E 2 size variants) and TTS/STT routes still fall through to read-time enforcement; those are follow-ups. --- .../spend_tracking/budget_reservation.py | 41 +++++ .../proxy/test_budget_reservation.py | 148 ++++++++++++++++++ 2 files changed, 189 insertions(+) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index feae368d23..1b11af84d0 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -827,6 +827,13 @@ def _estimate_request_max_cost_for_model( if model_info is None: return None + image_cost = _estimate_image_generation_cost( + request_body=request_body, + model_info=model_info, + ) + if image_cost is not None: + return image_cost + 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")) input_tokens = _estimate_input_tokens( @@ -858,6 +865,40 @@ def _estimate_request_max_cost_for_model( return cost +def _estimate_image_generation_cost( + request_body: dict, + model_info: Dict[str, Any], +) -> Optional[float]: + """ + Reserve `n × per-image cost` for image-generation requests so concurrent + requests against a depleted budget cannot all slip past the admission gate + onto the provider. Token-based pricing (e.g. gpt-image-1) is handled by + the chat-route token path; per-pixel and size/quality-tiered pricing + (DALL-E 2 size variants, premium tiers) are not handled here and fall + through to read-time enforcement. + + The "output" vs "input" cost-per-image naming is inconsistent across + providers — OpenAI's dall-e-3 entry uses ``input_cost_per_image`` while + aiml/dall-e-3 uses ``output_cost_per_image`` — so both are summed. + """ + output_cost_per_image = _to_float(model_info.get("output_cost_per_image")) + input_cost_per_image = _to_float(model_info.get("input_cost_per_image")) + is_image_gen = ( + model_info.get("mode") == "image_generation" + or output_cost_per_image is not None + or input_cost_per_image is not None + ) + if not is_image_gen: + return None + + cost_per_image = (output_cost_per_image or 0.0) + (input_cost_per_image or 0.0) + if cost_per_image <= 0: + return None + + n = _to_int(request_body.get("n")) or 1 + return cost_per_image * max(n, 1) + + def _get_model_cost_info( model: str, llm_router: Optional[Router], diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 0e2ca98a11..751e58e742 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -699,6 +699,154 @@ async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequest await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_should_reserve_image_generation_cost_per_image( + spend_counter_state, +): + """Image-generation requests reserve `n × per-image cost` so concurrent + requests against a depleted budget cannot all bypass the admission gate. + The OpenAI ``dall-e-3`` entry exposes the per-image price as + ``input_cost_per_image`` (a naming quirk), while other providers use + ``output_cost_per_image`` — both must be honored.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-image-gen", + spend=0.0, + max_budget=10.0, + ) + await key_cache.async_set_cache(key="key-image-gen", value=valid_token) + + request_body = {"model": "dall-e-3", "prompt": "a cat", "n": 3} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "mode": "image_generation", + "input_cost_per_image": 0.04, + }, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/v1/images/generations", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.12) # 3 × $0.04 + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_reject_concurrent_image_request_against_depleted_budget( + spend_counter_state, +): + """Greptile P1 regression: with image-gen reservation in place, a second + concurrent image request against a budget already pinned at the cap by + the first reservation must raise BudgetExceededError instead of + silently reaching the provider.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-image-deplete", + spend=0.0, + team_id="team-image-deplete", + ) + team_object = LiteLLM_TeamTable( + team_id="team-image-deplete", + max_budget=0.04, + spend=0.0, + ) + await key_cache.async_set_cache( + key=f"team_id:{team_object.team_id}", + value=team_object, + ) + + request_body = {"model": "dall-e-3", "prompt": "a cat"} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "mode": "image_generation", + "input_cost_per_image": 0.04, + }, + ): + first = await reserve_budget_for_request( + request_body=request_body, + route="/v1/images/generations", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert first is not None + + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=request_body, + route="/v1/images/generations", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + await release_budget_reservation(first) + + +@pytest.mark.asyncio +async def test_should_skip_reservation_for_per_pixel_image_model( + spend_counter_state, +): + """DALL-E 2-style per-pixel pricing depends on the requested ``size``, + which we don't decode here. Fall through to read-time enforcement + rather than guess.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-image-per-pixel", + spend=0.0, + max_budget=1.0, + ) + await key_cache.async_set_cache(key="key-image-per-pixel", value=valid_token) + + request_body = {"model": "dall-e-2", "prompt": "a cat", "size": "256x256"} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "mode": "image_generation", + "input_cost_per_pixel": 2.4414e-07, + "output_cost_per_pixel": 0.0, + }, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/v1/images/generations", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is None + + def test_should_start_window_without_reset_at_at_duration_boundary(): before = datetime.now(timezone.utc) - timedelta(hours=1) From 963cb4694dd39f43d01e9d02da03784dbb0b5fb7 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 9 May 2026 09:16:27 -0700 Subject: [PATCH 3/3] fix(proxy): gate image-gen reservation strictly on model mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous detection treated any model with input_cost_per_image or output_cost_per_image as image generation. Several chat and embedding models carry those fields to price multimodal vision input, not generated images: - gemini-3.1-pro-preview (mode=chat) has output_cost_per_image=0.00012 alongside input/output token pricing. - azure/gpt-realtime-* (mode=chat) has input_cost_per_image=5e-6. - amazon.titan-embed-image-v1 (mode=embedding) has input_cost_per_image=6e-5. For these models the image-gen branch fired first and reserved a fraction of a cent per request, short-circuiting the token-priced path entirely. Long Gemini chats reserved 1 × $0.00012 instead of the true token cost. Gate strictly on mode in {"image_generation", "image_edit"}. All 197 real image_generation entries and all 31 image_edit entries (Flux Kontext, Stability inpaint/outpaint, etc.) carry the right mode, so the field-presence fallback was unnecessary. Adds regression tests for the chat-model-with-image-cost-field case and for image_edit reservation. --- .../spend_tracking/budget_reservation.py | 20 ++-- .../proxy/test_budget_reservation.py | 104 ++++++++++++++++++ 2 files changed, 116 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 1b11af84d0..200a17e236 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -881,16 +881,20 @@ def _estimate_image_generation_cost( providers — OpenAI's dall-e-3 entry uses ``input_cost_per_image`` while aiml/dall-e-3 uses ``output_cost_per_image`` — so both are summed. """ - output_cost_per_image = _to_float(model_info.get("output_cost_per_image")) - input_cost_per_image = _to_float(model_info.get("input_cost_per_image")) - is_image_gen = ( - model_info.get("mode") == "image_generation" - or output_cost_per_image is not None - or input_cost_per_image is not None - ) - if not is_image_gen: + # Gate strictly on `mode`. Several chat and embedding models carry + # ``input_cost_per_image`` / ``output_cost_per_image`` to price multimodal + # *vision input* (e.g. ``gemini-3.1-pro-preview``, ``azure/gpt-realtime-*``, + # ``amazon.titan-embed-image-v1``). Falling back to "treat as image-gen if + # an image cost field is present" would short-circuit the token-priced + # path for those models and reserve a fraction of a cent instead of the + # true per-token cost. All real image-generation entries in + # ``model_prices_and_context_window.json`` carry ``mode: image_generation`` + # or ``mode: image_edit``, so the field-presence fallback is unnecessary. + if model_info.get("mode") not in ("image_generation", "image_edit"): return None + output_cost_per_image = _to_float(model_info.get("output_cost_per_image")) + input_cost_per_image = _to_float(model_info.get("input_cost_per_image")) cost_per_image = (output_cost_per_image or 0.0) + (input_cost_per_image or 0.0) if cost_per_image <= 0: return None diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 751e58e742..aa0f8d6327 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -847,6 +847,110 @@ async def test_should_skip_reservation_for_per_pixel_image_model( assert reservation is None +@pytest.mark.asyncio +async def test_should_use_token_pricing_for_chat_model_with_image_cost_field( + spend_counter_state, +): + """Several chat and embedding models carry ``input_cost_per_image`` / + ``output_cost_per_image`` to price multimodal vision *input*, not image + generation (e.g. gemini-3.1-pro-preview, azure/gpt-realtime-*, + amazon.titan-embed-image-v1). _estimate_image_generation_cost must gate + on ``mode`` so these models still go through the token-priced path — + otherwise a long chat reserves a fraction of a cent instead of the true + token cost.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-multimodal-chat", + spend=0.0, + max_budget=10.0, + ) + await key_cache.async_set_cache(key="key-multimodal-chat", value=valid_token) + + # Roughly the gemini-3.1-pro-preview shape: chat-mode model that + # carries an output_cost_per_image alongside token pricing. + output_cost_per_token = 1.2e-5 + request_body = { + "model": "gemini-3.1-pro-preview", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1000, + } + expected_cost = 1000 * output_cost_per_token # token-priced path, not 1 × $0.00012 + + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": output_cost_per_token, + "output_cost_per_image": 0.00012, + "max_output_tokens": 64000, + }, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + # Token-priced path: reservation ≈ output_tokens × output_cost_per_token, + # plus a small input-token contribution. Must NOT collapse to the + # per-image price ($0.00012) which would indicate the image-gen branch + # incorrectly fired for this chat model. + assert reservation["reserved_cost"] == pytest.approx(expected_cost, rel=0.05) + assert reservation["reserved_cost"] > 0.001 # well above per-image price + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_reserve_image_edit_cost_per_image( + spend_counter_state, +): + """``image_edit`` models (Flux Kontext, Stability inpaint/outpaint, etc.) + bill per generated image just like ``image_generation`` and must get + the same atomic per-image reservation.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-image-edit", + spend=0.0, + max_budget=10.0, + ) + await key_cache.async_set_cache(key="key-image-edit", value=valid_token) + + request_body = {"model": "stability/inpaint", "prompt": "a cat", "n": 2} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "mode": "image_edit", + "output_cost_per_image": 0.05, + }, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/v1/images/edits", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.10) # 2 × $0.05 + await release_budget_reservation(reservation) + + def test_should_start_window_without_reset_at_at_duration_boundary(): before = datetime.now(timezone.utc) - timedelta(hours=1)