mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
fix(proxy): narrow model-discovery budget bypass to explicit route set (#27923)
This commit is contained in:
Armaan Sandhu
parent
9491e0be49
commit
5a31f62d32
3 changed files with 125 additions and 203 deletions
|
|
@ -491,6 +491,20 @@ async def check_tools_allowlist(
|
|||
)
|
||||
|
||||
|
||||
# Read-only discovery routes that incur no spend. Kept narrower than info_routes so an exhausted
|
||||
# budget cannot reach side-effectful routes like /health/services (Slack/email/webhook). See #27923.
|
||||
MODEL_DISCOVERY_ROUTES = frozenset(
|
||||
{
|
||||
"/v1/models",
|
||||
"/models",
|
||||
"/model/info",
|
||||
"/v1/model/info",
|
||||
"/v2/model/info",
|
||||
"/model_group/info",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def common_checks( # noqa: PLR0915
|
||||
request_body: dict,
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
|
|
@ -534,13 +548,7 @@ async def common_checks( # noqa: PLR0915
|
|||
request_query_params=_safe_get_request_query_params(request=request),
|
||||
)
|
||||
|
||||
# Model-discovery / info routes (e.g. GET /v1/models, /models, /model/info) perform no
|
||||
# inference and incur no spend, so a budget must never block them. `_global_proxy_budget_check`
|
||||
# already exempts /v1/models and /models; mirror that here so the entity-level budget checks
|
||||
# (team / key / org / user / end-user / project) stay consistent and OpenAI-compatible clients
|
||||
# (Open WebUI, Cursor, Aider, etc.) can still discover models after a budget is exhausted.
|
||||
# See https://github.com/BerriAI/litellm/issues/27923
|
||||
if RouteChecks.is_info_route(route=route):
|
||||
if route in MODEL_DISCOVERY_ROUTES:
|
||||
skip_budget_checks = True
|
||||
|
||||
# 1. If team is blocked
|
||||
|
|
|
|||
|
|
@ -1,195 +0,0 @@
|
|||
"""
|
||||
Tests for budget bypass on model-discovery / info routes.
|
||||
|
||||
Model-discovery and info endpoints (e.g. GET /v1/models, /models, /model/info)
|
||||
perform no inference and incur no spend, so an exhausted budget must never block
|
||||
them. Otherwise OpenAI-compatible clients (Open WebUI, Cursor, Aider, Continue,
|
||||
LibreChat, ...) that call GET /v1/models for model discovery break entirely once
|
||||
any team/key/org/user budget is exceeded.
|
||||
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/27923
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
||||
# Routes that must never be blocked by a budget. Sourced from
|
||||
# LiteLLMRoutes.info_routes (the same group RouteChecks.is_info_route checks).
|
||||
MODEL_DISCOVERY_ROUTES = [
|
||||
"/v1/models",
|
||||
"/models",
|
||||
"/model/info",
|
||||
"/v1/model/info",
|
||||
"/v2/model/info",
|
||||
"/model_group/info",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_router():
|
||||
"""A router with a single paid (non-zero-cost) model."""
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "cloud-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "sk-test",
|
||||
},
|
||||
"model_info": {"id": "cloud-model-id"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_proxy_logging():
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=None)
|
||||
|
||||
async def mock_budget_alerts(*args, **kwargs):
|
||||
pass
|
||||
|
||||
proxy_logging.budget_alerts = mock_budget_alerts
|
||||
return proxy_logging
|
||||
|
||||
|
||||
class TestTeamBudgetInfoRouteBypass:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", MODEL_DISCOVERY_ROUTES)
|
||||
async def test_team_over_budget_can_list_models(
|
||||
self, route, mock_router, mock_proxy_logging
|
||||
):
|
||||
"""A team that is over budget can still reach model-discovery routes."""
|
||||
team_object = LiteLLM_TeamTable(
|
||||
team_id="test-team", spend=150.0, max_budget=100.0
|
||||
)
|
||||
|
||||
result = await common_checks(
|
||||
request_body={},
|
||||
team_object=team_object,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route=route,
|
||||
llm_router=mock_router,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_over_budget_still_blocked_on_inference(
|
||||
self, mock_router, mock_proxy_logging
|
||||
):
|
||||
"""Control: inference routes are still budget-enforced (no regression)."""
|
||||
team_object = LiteLLM_TeamTable(
|
||||
team_id="test-team", spend=150.0, max_budget=100.0
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await common_checks(
|
||||
request_body={"model": "cloud-model"},
|
||||
team_object=team_object,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=mock_router,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
class TestUserBudgetInfoRouteBypass:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", MODEL_DISCOVERY_ROUTES)
|
||||
async def test_user_over_budget_can_list_models(
|
||||
self, route, mock_router, mock_proxy_logging
|
||||
):
|
||||
user_object = LiteLLM_UserTable(
|
||||
user_id="test-user", spend=100.0, max_budget=50.0
|
||||
)
|
||||
|
||||
result = await common_checks(
|
||||
request_body={},
|
||||
team_object=None,
|
||||
user_object=user_object,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route=route,
|
||||
llm_router=mock_router,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
valid_token=UserAPIKeyAuth(token="test-token", user_id="test-user"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestEndUserBudgetInfoRouteBypass:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", MODEL_DISCOVERY_ROUTES)
|
||||
async def test_end_user_over_budget_can_list_models(
|
||||
self, route, mock_router, mock_proxy_logging
|
||||
):
|
||||
end_user_object = LiteLLM_EndUserTable(
|
||||
user_id="end-user-123",
|
||||
spend=50.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=20.0),
|
||||
blocked=False,
|
||||
)
|
||||
|
||||
result = await common_checks(
|
||||
request_body={"user": "end-user-123"},
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=end_user_object,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route=route,
|
||||
llm_router=mock_router,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
valid_token=UserAPIKeyAuth(token="test-token"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestGlobalProxyBudgetInfoRouteBypass:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", MODEL_DISCOVERY_ROUTES)
|
||||
async def test_global_proxy_over_budget_can_list_models(
|
||||
self, route, mock_router, mock_proxy_logging
|
||||
):
|
||||
with patch.object(litellm, "max_budget", 100.0):
|
||||
result = await common_checks(
|
||||
request_body={},
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=150.0,
|
||||
general_settings={},
|
||||
route=route,
|
||||
llm_router=mock_router,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
valid_token=UserAPIKeyAuth(token="test-token"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
assert result is True
|
||||
|
|
@ -1628,7 +1628,8 @@ async def test_reject_clientside_metadata_tags_non_llm_route():
|
|||
@pytest.mark.asyncio
|
||||
async def test_reject_clientside_metadata_tags_allows_key_tags_without_client_tags():
|
||||
"""Key metadata.tags are injected after the reject check; requests without
|
||||
client metadata.tags must not be blocked when reject_clientside_metadata_tags is on."""
|
||||
client metadata.tags must not be blocked when reject_clientside_metadata_tags is on.
|
||||
"""
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
|
@ -3513,3 +3514,111 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias():
|
|||
for c in cache2.async_set_cache.await_args_list
|
||||
]
|
||||
assert written_keys_aliasless == ["team_id:team-no-alias"]
|
||||
|
||||
|
||||
MODEL_DISCOVERY_ROUTES = [
|
||||
"/v1/models",
|
||||
"/models",
|
||||
"/model/info",
|
||||
"/v1/model/info",
|
||||
"/v2/model/info",
|
||||
"/model_group/info",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", MODEL_DISCOVERY_ROUTES)
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_discovery_route_bypasses_team_budget(route):
|
||||
"""Regression for #27923: an exhausted team budget must not block model-discovery routes,
|
||||
otherwise OpenAI-compatible clients calling GET /v1/models at startup break."""
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
|
||||
|
||||
result = await common_checks(
|
||||
request_body={},
|
||||
team_object=team_object,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route=route,
|
||||
llm_router=None,
|
||||
proxy_logging_obj=AsyncMock(),
|
||||
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_discovery_route_bypasses_user_budget():
|
||||
"""Regression for #27923: an exhausted user budget must not block model discovery."""
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
user_object = LiteLLM_UserTable(user_id="test-user", spend=100.0, max_budget=50.0)
|
||||
|
||||
result = await common_checks(
|
||||
request_body={},
|
||||
team_object=None,
|
||||
user_object=user_object,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/v1/models",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=AsyncMock(),
|
||||
valid_token=UserAPIKeyAuth(token="test-token", user_id="test-user"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
assert result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_side_effectful_info_route_still_enforces_budget():
|
||||
"""#27923 keeps the bypass narrow: /health/services can fire Slack/email/webhook test
|
||||
messages, so an exhausted budget must still block it. Widening the exemption back to
|
||||
is_info_route() would regress this."""
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await common_checks(
|
||||
request_body={},
|
||||
team_object=team_object,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/health/services",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=AsyncMock(),
|
||||
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inference_route_still_enforces_team_budget():
|
||||
"""Control for #27923: inference routes stay fully budget-enforced."""
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await common_checks(
|
||||
request_body={},
|
||||
team_object=team_object,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=AsyncMock(),
|
||||
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue