From 1416ab292d0d4a2952eabec3c3116ad40206f3bf Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 20:09:09 +0000 Subject: [PATCH] fix(proxy): enforce blocked customers and fix /customer/block 500 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 5 + litellm/proxy/auth/auth_checks.py | 10 ++ .../customer_endpoints.py | 93 ++++++++----------- .../proxy/auth/test_auth_checks.py | 50 ++++++++++ .../test_customer_endpoints.py | 81 ++++++++++++++++ 5 files changed, 183 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b7dc3a7c73..cf17436da1b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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"] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3166e58106c..7ef2c38d0a2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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"): diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 9446d903441..f7f2e9647cb 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -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: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index f2ab1bc965b..45231c9ec5f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 9b17d8ccc7d..afc6548d93f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -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"]