From d720b3d369eca916c87cd044acc6d2b192b8519f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 8 Jul 2025 16:46:05 -0700 Subject: [PATCH] [Bug fix] Multiple API Keys Created on Startup When max_budget is Enabled (#12436) * fix _add_proxy_budget_to_db * fix if table_name is not None and table_name == "user": * revert earlier change * test_add_proxy_budget_to_db_only_creates_user_no_keys * lint fix --- .../proxy/_experimental/out/onboarding.html | 1 - .../key_management_endpoints.py | 12 ++--- litellm/proxy/proxy_server.py | 1 + tests/test_litellm/proxy/test_proxy_server.py | 50 +++++++++++++++++++ 4 files changed, 54 insertions(+), 10 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/onboarding.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 061a432ac3..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 13f0d1753d..e9097f4c2f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1497,11 +1497,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 ] = None, # object_permission_id <-> LiteLLM_ObjectPermissionTable object_permission: Optional[LiteLLM_ObjectPermissionBase] = None, ): - from litellm.proxy.proxy_server import ( - litellm_proxy_budget_name, - premium_user, - prisma_client, - ) + from litellm.proxy.proxy_server import premium_user, prisma_client if prisma_client is None: raise Exception( @@ -1660,10 +1656,8 @@ async def generate_key_helper_fn( # noqa: PLR0915 table_name="user", update_key_values=update_key_values, ) - if user_id == litellm_proxy_budget_name or ( - table_name is not None and table_name == "user" - ): - # do not create a key for litellm_proxy_budget_name or if table name is set to just 'user' + if table_name is not None and table_name == "user": + # do not create a key if table name is set to just 'user' # we only need to ensure this exists in the user table # the LiteLLM_VerificationToken table will increase in size if we don't do this check return user_data diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3e15a2ecf5..eaef3c0887 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3305,6 +3305,7 @@ class ProxyStartupEvent: asyncio.create_task( generate_key_helper_fn( # type: ignore request_type="user", + table_name="user", user_id=litellm_proxy_budget_name, duration=None, models=[], diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 75fffeac5a..d33b637f8e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -611,3 +611,53 @@ async def test_get_config_from_file(tmp_path, monkeypatch): result = await proxy_config._get_config_from_file(None) assert result == test_config + + +@pytest.mark.asyncio +async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): + """ + Test that _add_proxy_budget_to_db only creates a user and no keys are added. + + This validates that generate_key_helper_fn is called with table_name="user" + which should prevent key creation in LiteLLM_VerificationToken table. + """ + from unittest.mock import AsyncMock, patch + + import litellm + from litellm.proxy.proxy_server import ProxyStartupEvent + + # Set up required litellm settings + litellm.budget_duration = "30d" + litellm.max_budget = 100.0 + + litellm_proxy_budget_name = "litellm-proxy-budget" + + # Mock generate_key_helper_fn to capture its call arguments + mock_generate_key_helper = AsyncMock(return_value={ + "user_id": litellm_proxy_budget_name, + "max_budget": 100.0, + "budget_duration": "30d", + "spend": 0, + "models": [], + }) + + # Patch generate_key_helper_fn in proxy_server where it's being called from + with patch("litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper): + # Call the function under test + ProxyStartupEvent._add_proxy_budget_to_db(litellm_proxy_budget_name) + + # Allow async task to complete + import asyncio + await asyncio.sleep(0.1) + + # Verify that generate_key_helper_fn was called + mock_generate_key_helper.assert_called_once() + call_args = mock_generate_key_helper.call_args + + # Verify critical parameters that prevent key creation + assert call_args.kwargs["request_type"] == "user" + assert call_args.kwargs["table_name"] == "user" + assert call_args.kwargs["user_id"] == litellm_proxy_budget_name + assert call_args.kwargs["max_budget"] == 100.0 + assert call_args.kwargs["budget_duration"] == "30d" + assert call_args.kwargs["query_type"] == "update_data"