mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(proxy): enforce blocked customers and fix /customer/block 500
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo
parent
52b4037177
commit
1416ab292d
5 changed files with 183 additions and 56 deletions
|
|
@ -3636,6 +3636,11 @@ class ProxyErrorTypes(str, enum.Enum):
|
|||
Tool is not in the allowed tools list for this key/team
|
||||
"""
|
||||
|
||||
blocked_user = "blocked_user"
|
||||
"""
|
||||
End-user is blocked from making requests
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_model_access_error_type_for_object(
|
||||
cls, object_type: Literal["key", "user", "team", "org", "project"]
|
||||
|
|
|
|||
|
|
@ -554,6 +554,16 @@ async def common_checks(
|
|||
if team_object is not None and team_object.blocked is True:
|
||||
raise Exception(f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin.")
|
||||
|
||||
# 1.05 If end-user is blocked
|
||||
if end_user_object is not None and end_user_object.blocked is True:
|
||||
raise ProxyException(
|
||||
message=f"User blocked from making LLM API Calls. User={end_user_object.user_id}. "
|
||||
"Update via `/customer/unblock` if you're an admin.",
|
||||
type=ProxyErrorTypes.blocked_user,
|
||||
param="user",
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
# 2. If team can call model (or key's access_group_ids grant it)
|
||||
if _model and team_object:
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"):
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ All /customer management endpoints
|
|||
"""
|
||||
|
||||
#### END-USER/CUSTOMER MANAGEMENT ####
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Iterable, Sequence
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -52,6 +52,35 @@ async def _invalidate_cached_end_users(end_user_ids: Iterable[str]) -> None:
|
|||
await delete_cached_end_user_object(end_user_id=end_user_id, user_api_key_cache=user_api_key_cache)
|
||||
|
||||
|
||||
async def _set_end_users_blocked(user_ids: Sequence[str], blocked: bool) -> tuple[LiteLLM_EndUserTable, ...]:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value})
|
||||
|
||||
try:
|
||||
records: Final = tuple(
|
||||
LiteLLM_EndUserTable.model_validate(
|
||||
(
|
||||
await EndUserRepository(prisma_client).table.upsert(
|
||||
where={"user_id": user_id},
|
||||
data={
|
||||
"create": {"user_id": user_id, "blocked": blocked},
|
||||
"update": {"blocked": blocked},
|
||||
},
|
||||
)
|
||||
).model_dump()
|
||||
)
|
||||
for user_id in user_ids
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("An error occurred - %s", e)
|
||||
raise HTTPException(status_code=500, detail={"error": str(e)})
|
||||
|
||||
await _invalidate_cached_end_users(user_ids)
|
||||
return records
|
||||
|
||||
|
||||
def _to_customer_response(record: BaseModel) -> CustomerResponse:
|
||||
"""Validate a raw end-user DB row into the typed customer response.
|
||||
|
||||
|
|
@ -91,31 +120,7 @@ async def block_user(data: BlockUsers):
|
|||
}'
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
try:
|
||||
records: Final = []
|
||||
if prisma_client is not None:
|
||||
for id in data.user_ids:
|
||||
record = await EndUserRepository(prisma_client).table.upsert(
|
||||
where={"user_id": id},
|
||||
data={
|
||||
"create": {"user_id": id, "blocked": True},
|
||||
"update": {"blocked": True},
|
||||
},
|
||||
)
|
||||
records.append(record)
|
||||
await _invalidate_cached_end_users(data.user_ids)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": "Postgres DB Not connected"},
|
||||
)
|
||||
|
||||
return {"blocked_users": records}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("An error occurred - %s", e)
|
||||
raise HTTPException(status_code=500, detail={"error": str(e)})
|
||||
return BlockUsersResponse(blocked_users=list(await _set_end_users_blocked(user_ids=data.user_ids, blocked=True)))
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -130,7 +135,7 @@ async def block_user(data: BlockUsers):
|
|||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=UnblockUsersResponse,
|
||||
)
|
||||
async def unblock_user(data: BlockUsers):
|
||||
async def unblock_user(data: BlockUsers) -> UnblockUsersResponse:
|
||||
"""
|
||||
[BETA] Unblock calls with this user id
|
||||
|
||||
|
|
@ -143,38 +148,14 @@ async def unblock_user(data: BlockUsers):
|
|||
}'
|
||||
```
|
||||
"""
|
||||
try:
|
||||
from enterprise.enterprise_hooks.blocked_user_list import (
|
||||
_ENTERPRISE_BlockedUserList,
|
||||
)
|
||||
except ImportError:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Blocked user check was never set. This call has no effect."
|
||||
+ CommonProxyErrors.missing_enterprise_package_docker.value
|
||||
},
|
||||
)
|
||||
|
||||
if (
|
||||
not any(isinstance(x, _ENTERPRISE_BlockedUserList) for x in litellm.callbacks)
|
||||
or litellm.blocked_user_list is None
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Blocked user check was never set. This call has no effect."},
|
||||
)
|
||||
await _set_end_users_blocked(user_ids=data.user_ids, blocked=False)
|
||||
|
||||
if isinstance(litellm.blocked_user_list, list):
|
||||
for id in data.user_ids:
|
||||
litellm.blocked_user_list.remove(id)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": "`blocked_user_list` must be set as a list. Filepaths can't be updated."},
|
||||
)
|
||||
remaining: Final = tuple(user_id for user_id in litellm.blocked_user_list if user_id not in set(data.user_ids))
|
||||
litellm.blocked_user_list = list(remaining)
|
||||
return UnblockUsersResponse(blocked_users=list(remaining))
|
||||
|
||||
return {"blocked_users": litellm.blocked_user_list}
|
||||
return UnblockUsersResponse(blocked_users=[])
|
||||
|
||||
|
||||
def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None:
|
||||
|
|
|
|||
|
|
@ -5723,3 +5723,53 @@ async def test_delete_cached_end_user_object_still_broadcasts_when_eviction_fail
|
|||
("end_user_id:eu-1", "user_api_key"),
|
||||
("end_user_validation:eu-1", "user_api_key"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_rejects_blocked_end_user():
|
||||
"""A customer blocked via /customer/block must not be able to make LLM calls.
|
||||
|
||||
Enforcement used to live only in the enterprise blocked_user_list hook, so on a default
|
||||
install `blocked: true` in the DB had no effect at all.
|
||||
"""
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await common_checks(
|
||||
request_body={"user": "blocked-customer"},
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=LiteLLM_EndUserTable(user_id="blocked-customer", blocked=True),
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=AsyncMock(),
|
||||
valid_token=UserAPIKeyAuth(token="test-token"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.blocked_user
|
||||
assert "blocked-customer" in exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_allows_unblocked_end_user():
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
result = await common_checks(
|
||||
request_body={"user": "ok-customer"},
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=LiteLLM_EndUserTable(user_id="ok-customer", blocked=False),
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/v1/chat/completions",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=AsyncMock(),
|
||||
valid_token=UserAPIKeyAuth(token="test-token"),
|
||||
request=MagicMock(),
|
||||
)
|
||||
|
||||
assert result is True
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from fastapi import FastAPI, HTTPException, Request, status
|
|||
from fastapi.responses import JSONResponse
|
||||
from fastapi.routing import APIRoute
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_EndUserTable,
|
||||
|
|
@ -860,3 +861,83 @@ def test_new_customer_evicts_cached_negative_validation(mock_prisma_client, mock
|
|||
|
||||
assert response.status_code == 200
|
||||
assert cache.in_memory_cache.get_cache("end_user_validation:c1") is None
|
||||
|
||||
|
||||
class _PrismaEndUserRow(BaseModel):
|
||||
"""Stand-in for the prisma-generated row model, which is a distinct class from
|
||||
litellm's LiteLLM_EndUserTable and has no dict-like `.get`."""
|
||||
|
||||
user_id: str
|
||||
blocked: bool
|
||||
spend: float = 0.0
|
||||
|
||||
|
||||
def test_block_customer_returns_200_for_prisma_row(mock_prisma_client, mock_user_api_key_auth):
|
||||
"""The prisma row is not a LiteLLM_EndUserTable, so returning it raw made
|
||||
response_model validation blow up with a 500 on every /customer/block call."""
|
||||
mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(
|
||||
return_value=_PrismaEndUserRow(user_id="blocked-1", blocked=True)
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/customer/block",
|
||||
json={"user_ids": ["blocked-1"]},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["blocked_users"][0] == {
|
||||
"user_id": "blocked-1",
|
||||
"blocked": True,
|
||||
"alias": None,
|
||||
"spend": 0.0,
|
||||
"allowed_model_region": None,
|
||||
"default_model": None,
|
||||
"budget_id": None,
|
||||
"litellm_budget_table": None,
|
||||
"object_permission_id": None,
|
||||
"object_permission": None,
|
||||
}
|
||||
|
||||
|
||||
def test_unblock_customer_clears_blocked_flag_in_db(mock_prisma_client, mock_user_api_key_auth):
|
||||
"""Unblocking used to only mutate the in-memory litellm.blocked_user_list, so a customer
|
||||
blocked through /customer/block stayed blocked in the DB forever."""
|
||||
upsert = AsyncMock(return_value=_PrismaEndUserRow(user_id="blocked-1", blocked=False))
|
||||
mock_prisma_client.db.litellm_endusertable.upsert = upsert
|
||||
|
||||
response = client.post(
|
||||
"/customer/unblock",
|
||||
json={"user_ids": ["blocked-1"]},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert upsert.call_args.kwargs["data"]["update"] == {"blocked": False}
|
||||
|
||||
|
||||
def test_block_and_unblock_invalidate_cached_end_users(mock_prisma_client, mock_user_api_key_auth, monkeypatch):
|
||||
invalidated: list[str] = []
|
||||
|
||||
async def _record(end_user_id: str, user_api_key_cache) -> None:
|
||||
invalidated.append(end_user_id)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.customer_endpoints.delete_cached_end_user_object",
|
||||
_record,
|
||||
)
|
||||
mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(
|
||||
return_value=_PrismaEndUserRow(user_id="blocked-1", blocked=True)
|
||||
)
|
||||
|
||||
for route in ("/customer/block", "/customer/unblock"):
|
||||
assert (
|
||||
client.post(
|
||||
route,
|
||||
json={"user_ids": ["blocked-1"]},
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
|
||||
assert invalidated == ["blocked-1", "blocked-1"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue