diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 1d296611bf..200a17e236 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, @@ -855,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( @@ -886,6 +865,44 @@ 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. + """ + # 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 + + 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], @@ -946,6 +963,9 @@ def _estimate_input_tokens( return None +DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK = 16384 + + def _estimate_output_tokens( request_body: dict, route: str, @@ -954,15 +974,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..aa0f8d6327 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,66 +695,91 @@ 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( +async def test_should_reserve_image_generation_cost_per_image( spend_counter_state, - monkeypatch, ): + """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-budget-double-resize", - spend=0.2, - max_budget=1.0, - team_id="team-budget-double-resize", + 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-budget-double-resize", - spend=0.2, - max_budget=1.0, + team_id="team-image-deplete", + max_budget=0.04, + spend=0.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, + 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.estimate_request_max_cost", - return_value=None, + "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( + first = await reserve_budget_for_request( request_body=request_body, - route="/chat/completions", + route="/v1/images/generations", llm_router=None, valid_token=valid_token, team_object=team_object, @@ -754,32 +788,167 @@ async def test_should_shrink_uncapped_reservation_multiple_times( 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 + + +@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 - 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) - + # 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) - 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) + +@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(): @@ -1047,62 +1216,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 +1605,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)