From 4e20fd138201d8939fd7e1be8c337f606ccd1dbe Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Wed, 22 Jul 2026 23:09:49 +0530 Subject: [PATCH] fix(proxy): fix double refund and multimodal undercounting in project io token reservation --- .../hooks/parallel_request_limiter_v3.py | 55 +++++++- .../proxy/hooks/test_tpm_concurrent.py | 120 ++++++++++++++++++ 2 files changed, 168 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7d0f6d2c57e..502a93404f7 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -2542,6 +2542,41 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): llm_provider=llm_provider, ) + def _estimate_precise_input_tokens(self, data: dict, model: str | None) -> int: + """ + Model-aware input token estimate for the project ITPM reservation, + using ``litellm.token_counter`` -- the same approach the + deployment-level itpm/otpm check uses in + ``io_token_rate_limit_check.py``. Unlike the cheap char-count + estimate the combined-TPM path uses, this accounts for image/tool + content so a burst of multimodal or tool-heavy requests can't each + reserve only the one-token floor and blow past ITPM before post-call + reconciliation catches up. + + Falls back to the cheap char-count estimate if ``token_counter`` + can't resolve a tokenizer for this model (e.g. an unrecognized + custom model name) or otherwise raises. + """ + from litellm import token_counter + + try: + return max( + 0, + int( + token_counter( + model=model or "", + messages=data.get("messages"), + text=data.get("prompt") or data.get("input"), + tools=data.get("tools"), + tool_choice=data.get("tool_choice"), + use_default_image_token_count=True, + ) + ), + ) + except Exception: # noqa: BLE001 - any tokenizer/model-resolution failure degrades to the cheap estimate, never a 500 + estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data) + return estimated_input_tokens + async def _reserve_project_io_tokens_or_raise( self, descriptors: list[RateLimitDescriptor], @@ -2575,10 +2610,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] min_configured_otpm_limit = min(configured_otpm_limits) if configured_otpm_limits else None - estimated_input_tokens, estimated_output_tokens = self._estimate_input_and_output_tokens( + _, estimated_output_tokens = self._estimate_input_and_output_tokens( data=data, min_configured_tpm_limit=min_configured_otpm_limit, ) + estimated_input_tokens = self._estimate_precise_input_tokens(data=data, model=requested_model) estimated_input_tokens = max(estimated_input_tokens, 1) estimated_output_tokens = max(estimated_output_tokens, 1) @@ -2605,12 +2641,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if io_response["overall_code"] == "OVER_LIMIT": # A combined-TPM reservation may have already succeeded above for # this same request; refund it too, or its counter stays inflated - # until the window's TTL expires. - await self._refund_reserved_tokens( - scopes=tpm_reservation_scopes, - amount=tpm_reservation_amount, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) + # until the window's TTL expires. Mark it released so the + # ProxyRateLimitError we're about to raise doesn't get refunded + # a second time when async_post_call_failure_hook sees the same + # (still-stashed) reservation and refunds it again. + if tpm_reservation_amount > 0: + await self._refund_reserved_tokens( + scopes=tpm_reservation_scopes, + amount=tpm_reservation_amount, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + self._mark_reservation_released(data) acquisition = self._get_parallel_slot_acquisition(kwargs=data) if acquisition is not None: await self._release_parallel_request_slots( diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 3f6cbe239bc..a1a786f309c 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -1523,5 +1523,125 @@ async def test_proxy_rejection_refunds_itpm_otpm_only_reservation_with_no_combin assert otpm_after == 0, f"OTPM-only reservation leaked on proxy rejection: {otpm_after}" +@pytest.mark.asyncio +async def test_otpm_rejection_does_not_double_refund_combined_tpm(rate_limiter): + """ + Regression for a High-severity review finding: when the project ITPM + reservation succeeds but OTPM is then over limit, + _reserve_project_io_tokens_or_raise rolls back the combined-TPM + reservation that already succeeded earlier in the same pre-call, then + raises. If it doesn't also mark that reservation released, + async_post_call_failure_hook -- which fires next in the real request + lifecycle, since raising from async_pre_call_hook triggers it -- sees + the same still-stashed reservation and refunds it a second time, + driving the combined TPM counter negative and letting a caller push + past the project's real TPM budget by repeatedly triggering OTPM + rejections. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-double-refund") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-double-refund", + project_metadata={ + "model_tpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_itpm_limit": {"bedrock_mantle/claude-opus": 100000}, + "model_otpm_limit": {"bedrock_mantle/claude-opus": 5}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [{"role": "user", "content": "hello there, this is a test message"}], + "max_tokens": 60, # blows past the 5-token OTPM limit + } + + tpm_counter_key = handler.create_rate_limit_keys( + key="model_per_project", value="proj-double-refund:bedrock_mantle/claude-opus", rate_limit_type="tokens" + ) + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429 + + tpm_after_pre_call = int(await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0) + assert tpm_after_pre_call == 0, f"combined TPM reservation not rolled back: {tpm_after_pre_call}" + + # In the real request lifecycle, async_post_call_failure_hook fires next + # for a pre-call rejection. It must not refund the same reservation again. + await handler.async_post_call_failure_hook( + request_data=data, + original_exception=exc_info.value, + user_api_key_dict=user_api_key_dict, + ) + + tpm_after_failure_hook = int(await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0) + assert tpm_after_failure_hook == 0, ( + f"combined TPM counter went negative from a double refund: {tpm_after_failure_hook}" + ) + + +@pytest.mark.asyncio +async def test_itpm_reservation_accounts_for_image_content_not_just_text(rate_limiter): + """ + Regression for a Medium-severity review finding: the ITPM reservation + used to estimate input tokens from message text length alone, so a + request with a tiny text prompt but real image content would reserve + almost nothing, letting a burst of such requests blow past the + configured ITPM limit before post-call usage reconciliation catches up. + _estimate_precise_input_tokens uses litellm.token_counter (with + use_default_image_token_count=True, so this test makes no network call) + to account for image content instead of only text length. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-multimodal-itpm") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-multimodal", + project_metadata={ + # Tighter than the ~250-token default image estimate, but far + # bigger than the handful of tokens the bare text "hi" would + # cost under the old char-count-only estimate. + "model_itpm_limit": {"bedrock_mantle/claude-opus": 50}, + }, + ) + + data = { + "model": "bedrock_mantle/claude-opus", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/some-image.png"}, + }, + ], + } + ], + } + + with pytest.raises(Exception) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + assert getattr(exc_info.value, "status_code", None) == 429, ( + "Expected the image content to push the ITPM reservation over the " + "50-token limit; if this doesn't raise, the estimate is undercounting " + "non-text content again." + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])