mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(auth): enforce key max_budget for proxy admin keys
The PROXY_ADMIN shortcut in the auth builder returned before Check 5, where non-admin keys enforce their own max_budget. An admin key with a configured max_budget therefore stayed usable after its spend exceeded the limit. Run the same virtual-key budget check on the admin path for LLM routes, keeping the existing zero-cost/MCP discovery skip conditions. Add a regression test for over-budget and under-budget admin keys. Signed-off-by: Anh Nguyen (Jerry) <92361793+anh-nguyenn@users.noreply.github.com>
This commit is contained in:
parent
dddbaa803d
commit
bf1f97291b
2 changed files with 110 additions and 0 deletions
|
|
@ -2664,6 +2664,21 @@ async def user_api_key_auth_builder(
|
|||
code=status.HTTP_401_UNAUTHORIZED,
|
||||
param=abbreviate_api_key(api_key=api_key),
|
||||
)
|
||||
# The admin shortcut below returns before Check 5, which is where
|
||||
# non-admin keys enforce their own max_budget. Run that check here so
|
||||
# a budgeted admin key is not usable after its spend limit is reached.
|
||||
if RouteChecks.is_llm_api_route(route=route) and not _should_skip_budget_checks(
|
||||
request_data=request_data,
|
||||
route=route,
|
||||
request=request,
|
||||
llm_router=llm_router,
|
||||
team_id=valid_token.team_id,
|
||||
):
|
||||
await virtual_key_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
valid_token = update_valid_token_with_end_user_params(
|
||||
valid_token=valid_token, end_user_params=end_user_params
|
||||
)
|
||||
|
|
|
|||
95
tests/unit/proxy/auth/test_proxy_admin_key_budget.py
Normal file
95
tests/unit/proxy/auth/test_proxy_admin_key_budget.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
from collections.abc import Iterator
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proxy_server_auth_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[UserAPIKeyAuth]:
|
||||
api_key = "sk-admin-key-budget-test"
|
||||
token = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
token=api_key,
|
||||
user_id="admin-user",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
max_budget=1.0,
|
||||
)
|
||||
|
||||
cache = AsyncMock()
|
||||
cache.async_get_cache = AsyncMock(return_value=None)
|
||||
cache.get_cache = AsyncMock(return_value=None)
|
||||
cache.delete_cache = MagicMock()
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.internal_usage_cache = MagicMock()
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.budget_alerts = AsyncMock(return_value=None)
|
||||
|
||||
for attr, value in {
|
||||
"prisma_client": MagicMock(),
|
||||
"user_api_key_cache": cache,
|
||||
"proxy_logging_obj": proxy_logging_obj,
|
||||
"master_key": "sk-master-key",
|
||||
"general_settings": {},
|
||||
"llm_model_list": [],
|
||||
"llm_router": None,
|
||||
"open_telemetry_logger": None,
|
||||
"model_max_budget_limiter": MagicMock(),
|
||||
"user_custom_auth": None,
|
||||
"jwt_handler": None,
|
||||
"litellm_proxy_admin_name": "admin",
|
||||
}.items():
|
||||
monkeypatch.setattr(ps, attr, value)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key",
|
||||
new_callable=AsyncMock,
|
||||
return_value=token,
|
||||
):
|
||||
yield token
|
||||
|
||||
|
||||
async def _authenticate(api_key: str) -> UserAPIKeyAuth:
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
return await user_api_key_auth_builder(
|
||||
request=request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_admin_key_over_budget_is_rejected(
|
||||
proxy_server_auth_state: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Admin keys skipped the key max_budget check and stayed usable after spend exceeded it."""
|
||||
monkeypatch.setattr(ps, "get_current_spend", AsyncMock(return_value=2.0))
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _authenticate(proxy_server_auth_state.api_key or "")
|
||||
|
||||
assert "Budget has been exceeded" in str(exc_info.value.message)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_admin_key_under_budget_is_allowed(
|
||||
proxy_server_auth_state: UserAPIKeyAuth, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Positive control: the same admin path still authenticates while under budget."""
|
||||
monkeypatch.setattr(ps, "get_current_spend", AsyncMock(return_value=0.5))
|
||||
|
||||
result = await _authenticate(proxy_server_auth_state.api_key or "")
|
||||
|
||||
assert result.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
Loading…
Add table
Reference in a new issue