From ace963b7a8fe0046b99e084cda4082af08253103 Mon Sep 17 00:00:00 2001 From: Charan Rathore Date: Mon, 28 Sep 2026 14:22:41 +0530 Subject: [PATCH] fix(proxy): reject dynamic limiter requests on usage lookup failure --- litellm/proxy/hooks/dynamic_rate_limiter.py | 7 +- .../proxy/hooks/test_dynamic_rate_limiter.py | 65 +++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index f4eac6ae5ae..7d8f3dccb01 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -8,6 +8,8 @@ from collections.abc import Callable from datetime import datetime from typing import Final +from fastapi import HTTPException + import litellm from litellm import ModelResponse, Router from litellm._logging import verbose_proxy_logger @@ -177,7 +179,10 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): verbose_proxy_logger.exception( "litellm.proxy.hooks.dynamic_rate_limiter.py::check_available_usage: Exception occurred - %s", e ) - return None, None, None, None, None + raise HTTPException( + status_code=503, + detail="Dynamic rate limiter usage is unavailable; request rejected.", + ) from e async def async_pre_call_hook( self, diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter.py index b630ff1605c..403d5212f2e 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter.py @@ -1,8 +1,13 @@ +import asyncio from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.dynamic_rate_limiter import ( DynamicRateLimiterCache, _PROXY_DynamicRateLimitHandler, @@ -42,3 +47,63 @@ async def test_handler_threads_time_fn_to_internal_cache(): ) await handler.internal_usage_cache.async_set_cache_sadd(model="my-fake-model", value=["p1", "p2"]) assert await handler.internal_usage_cache.async_get_cache(model="my-fake-model") == 2 + + +def fake_handler(*, cache_error=None, tpm=100, rpm=10): + usage_cache = MagicMock() + usage_cache.async_get_cache = AsyncMock(side_effect=cache_error, return_value=None) + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=usage_cache) + router = MagicMock() + router.get_model_group_info.return_value = MagicMock(tpm=tpm, rpm=rpm) + router.get_model_group_usage = AsyncMock(return_value=(0, 0)) + handler.update_variables(router) + return handler, usage_cache, router + + +@pytest.mark.asyncio +async def test_cache_lookup_failure_rejects_before_router_call(): + handler, _, router = fake_handler(cache_error=ConnectionError("cache offline")) + with pytest.raises(HTTPException) as exc: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", metadata={}), + cache=DualCache(), + data={"model": "test-model"}, + call_type="completion", + ) + assert exc.value.status_code == 503 + assert "cache offline" not in str(exc.value.detail) + router.get_model_group_usage.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_zero_quota_still_raises_rate_limit(): + handler, _, router = fake_handler(tpm=0) + with pytest.raises(ProxyRateLimitError) as exc: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", metadata={}), + cache=DualCache(), + data={"model": "gpt-4o-mini"}, + call_type="completion", + ) + assert exc.value.status_code == 429 + router.get_model_group_usage.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_unlimited_capacity_still_admits_request(): + handler, _, router = fake_handler(tpm=None, rpm=None) + user = UserAPIKeyAuth(api_key="sk-test", metadata={}) + with patch("litellm.proxy.hooks.dynamic_rate_limiter.asyncio.create_task") as create_task: + assert await handler.async_pre_call_hook(user, DualCache(), {"model": "test-model"}, "completion") is None + create_task.assert_not_called() + router.get_model_group_usage.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_successful_lookup_keeps_active_project_accounting(): + handler, _, _ = fake_handler() + user = UserAPIKeyAuth(api_key="sk-test", token="token-123", metadata={}) + handler.internal_usage_cache.async_set_cache_sadd = AsyncMock() # type: ignore[method-assign] + await handler.async_pre_call_hook(user, DualCache(), {"model": "test-model"}, "completion") + await asyncio.sleep(0) + handler.internal_usage_cache.async_set_cache_sadd.assert_awaited_once_with(model="test-model", value=[user.token])