scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

Merge 45b55bd60b into eb0e3f8c18

Browse source
This commit is contained in:
ASHISH 2026-08-27 19:16:23 -05:00 • committed by GitHub
commit 2aefb6a557
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 147 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/proxy/proxy_server.py
View file

@ -1225,13 +1225,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
prompt_injection_detection_obj.update_environment(router=llm_router)
verbose_proxy_logger.debug("prisma_client: %s", prisma_client)
if prisma_client is not None and litellm.max_budget > 0:
ProxyStartupEvent._add_proxy_budget_to_db()
asyncio.create_task(
ProxyStartupEvent._warm_global_spend_cache(
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
if prisma_client is not None:
ProxyStartupEvent._sync_proxy_budget_state(
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
ProxyStartupEvent._warn_budget_without_db(
max_budget=litellm.max_budget,
@ -8872,6 +8869,28 @@ class ProxyStartupEvent:
litellm_jwtauth=litellm_jwtauth,
)
@classmethod
def _sync_proxy_budget_state(
cls,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> None:
if litellm.max_budget > 0:
cls._add_proxy_budget_to_db()
asyncio.create_task(
cls._warm_global_spend_cache(
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
)
else:
asyncio.create_task(
cls._clear_stale_proxy_budget_from_db(
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
)
@classmethod
def _add_proxy_budget_to_db(cls):
"""Adds a global proxy budget to db"""
@ -8880,6 +8899,29 @@ class ProxyStartupEvent:
asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill())
@classmethod
async def _clear_stale_proxy_budget_from_db(
cls,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> None:
try:
await UserRepository(prisma_client).table.update_many(
where={
"user_id": {"in": [LITELLM_PROXY_BUDGET_NAME, LITELLM_PROXY_ADMIN_NAME]},
"max_budget": {"not": None},
},
data={
"max_budget": None,
"budget_duration": None,
"budget_reset_at": None,
},
)
await user_api_key_cache.async_delete_cache(key=LITELLM_PROXY_BUDGET_NAME)
await user_api_key_cache.async_delete_cache(key=LITELLM_PROXY_ADMIN_NAME)
except Exception as e:
verbose_proxy_logger.warning("Failed to clear stale proxy budget rows: %s", e)
@classmethod
async def _upsert_proxy_budget_with_reset_at_backfill(cls) -> None:
"""

98
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -2934,6 +2934,104 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at():
assert backfill_call.kwargs["data"]["spend"] == 0
@pytest.mark.asyncio
async def test_sync_proxy_budget_state_clears_stale_budget_when_max_budget_unset():
"""
Regression test for https://github.com/BerriAI/litellm/issues/35680.
When litellm_settings.max_budget is removed from config (litellm.max_budget
resolves to 0), startup must clear any max_budget/budget_duration previously
synced onto the proxy budget rows, on both the current sync target
("litellm-proxy-budget") and the legacy sync target ("default_user_id").
Without this, master-key traffic (which resolves to default_user_id) keeps
getting 429'd against a budget that no longer exists in config.
"""
from litellm.proxy.proxy_server import ProxyStartupEvent
original_max_budget = litellm.max_budget
litellm.max_budget = 0.0
mock_prisma = MagicMock()
mock_prisma.db.litellm_usertable.update_many = AsyncMock(return_value={"count": 2})
mock_cache = MagicMock()
mock_cache.async_delete_cache = AsyncMock()
try:
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
ProxyStartupEvent._sync_proxy_budget_state(
user_api_key_cache=mock_cache,
prisma_client=mock_prisma,
)
await asyncio.sleep(0.1)
mock_prisma.db.litellm_usertable.update_many.assert_called_once()
clear_call = mock_prisma.db.litellm_usertable.update_many.call_args
assert set(clear_call.kwargs["where"]["user_id"]["in"]) == {
"litellm-proxy-budget",
"default_user_id",
}
assert clear_call.kwargs["where"]["max_budget"] == {"not": None}
assert clear_call.kwargs["data"] == {
"max_budget": None,
"budget_duration": None,
"budget_reset_at": None,
}
deleted_keys = {call.kwargs["key"] for call in mock_cache.async_delete_cache.call_args_list}
assert deleted_keys == {"litellm-proxy-budget", "default_user_id"}
finally:
litellm.max_budget = original_max_budget
@pytest.mark.asyncio
async def test_sync_proxy_budget_state_syncs_when_max_budget_configured():
"""
When litellm.max_budget is set, startup must sync the budget onto the db
(existing behavior) and must NOT run the stale-budget clear path.
"""
from litellm.proxy.proxy_server import ProxyStartupEvent
original_max_budget = litellm.max_budget
original_budget_duration = litellm.budget_duration
litellm.max_budget = 100.0
litellm.budget_duration = "30d"
mock_prisma = MagicMock()
mock_prisma.db.litellm_usertable.update_many = AsyncMock(return_value={"count": 0})
mock_cache = MagicMock()
mock_cache.async_delete_cache = AsyncMock()
mock_generate_key_helper = AsyncMock(
return_value={
"user_id": "litellm-proxy-budget",
"max_budget": 100.0,
"budget_duration": "30d",
"spend": 0,
"models": [],
}
)
try:
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.generate_key_helper_fn",
mock_generate_key_helper,
),
):
ProxyStartupEvent._sync_proxy_budget_state(
user_api_key_cache=mock_cache,
prisma_client=mock_prisma,
)
await asyncio.sleep(0.1)
mock_generate_key_helper.assert_called_once()
mock_cache.async_delete_cache.assert_not_called()
finally:
litellm.max_budget = original_max_budget
litellm.budget_duration = original_budget_duration
@pytest.mark.asyncio
async def test_custom_ui_sso_sign_in_handler_config_loading():
"""