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"