From 7039c8b2ccef0693eba872c3a10d4bbb162fbe6d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 2 Aug 2026 02:30:41 +0000 Subject: [PATCH] fix(proxy): enforce EndUser.blocked in core auth and persist customer unblock --- litellm/proxy/auth/auth_checks.py | 6 + .../customer_endpoints.py | 62 +++++++---- .../proxy/auth/test_auth_checks.py | 52 +++++++++ .../test_customer_endpoints.py | 105 ++++++++++++++++++ 4 files changed, 204 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f876b303510..c5817f9a380 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -515,6 +515,7 @@ async def common_checks( 1. If team is blocked 1.1. If project is blocked + 1.2. If end user is blocked 2. If team can call model 2.2 If project can call model 3. If team is in budget @@ -549,6 +550,11 @@ 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.") + if end_user_object is not None and end_user_object.blocked is True: + raise Exception( + f"End user={end_user_object.user_id} is blocked. Update via `/customer/unblock` if you're an admin." + ) + # 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 09977fdce40..992acdce644 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -42,6 +42,21 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import ( router = APIRouter() +async def _invalidate_end_user_cache(user_ids: list[str]) -> None: + """Drop every cached view of an end-user row so a block/unblock takes effect now. + + Core auth caches the row under ``end_user_id:{id}`` and the enterprise + blocked-user hook caches it under ``litellm:end_user_id:{id}`` (60s ttl); + both live in the proxy's ``user_api_key_cache``, so a warmed entry would + otherwise keep serving the pre-change ``blocked`` value. + """ + from litellm.proxy.proxy_server import user_api_key_cache + + for user_id in user_ids: + await user_api_key_cache.async_delete_cache(key=f"end_user_id:{user_id}") + await user_api_key_cache.async_delete_cache(key=f"litellm:end_user_id:{user_id}") + + def _to_customer_response(record: BaseModel) -> CustomerResponse: """Validate a raw end-user DB row into the typed customer response. @@ -101,6 +116,8 @@ async def block_user(data: BlockUsers): detail={"error": "Postgres DB Not connected"}, ) + await _invalidate_end_user_cache(user_ids=data.user_ids) + return {"blocked_users": records} except Exception as e: verbose_proxy_logger.error(f"An error occurred - {e!s}") @@ -132,38 +149,39 @@ async def unblock_user(data: BlockUsers): }' ``` """ - try: - from enterprise.enterprise_hooks.blocked_user_list import ( - _ENTERPRISE_BlockedUserList, + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Postgres DB Not connected"}, ) - except ImportError: + + if isinstance(litellm.blocked_user_list, str): raise HTTPException( status_code=400, detail={ - "error": "Blocked user check was never set. This call has no effect." - + CommonProxyErrors.missing_enterprise_package_docker.value + "error": "`blocked_user_list` is set to a filepath, which can't be updated. " + "Remove the user id from that file and restart the proxy." }, ) - 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."}, + try: + await EndUserRepository(prisma_client).table.update_many( + where={"user_id": {"in": data.user_ids}}, + data={"blocked": False}, ) + except Exception as e: + verbose_proxy_logger.error(f"An error occurred - {e!s}") + raise HTTPException(status_code=500, detail={"error": str(e)}) + await _invalidate_end_user_cache(user_ids=data.user_ids) + + unblocked_ids = set(data.user_ids) 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."}, - ) + litellm.blocked_user_list = [user_id for user_id in litellm.blocked_user_list if user_id not in unblocked_ids] - return {"blocked_users": litellm.blocked_user_list} + return {"blocked_users": litellm.blocked_user_list or []} def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: @@ -629,6 +647,8 @@ async def update_end_user( raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}") verbose_proxy_logger.debug(f"received response from updating prisma client. response={response}") + await _invalidate_end_user_cache(user_ids=[data.user_id]) + return _to_customer_response(response) else: raise ValueError(f"user_id is required, passed user_id = {data.user_id}") diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5f3b0f36b95..8bc29d90393 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5233,3 +5233,55 @@ async def test_get_project_object_db_fetch_returns_cached_obj(): assert isinstance(result, LiteLLM_ProjectTableCachedObj) assert result.project_id == "p-1" assert result.project_alias == "proj" + + +@pytest.mark.asyncio +async def test_common_checks_rejects_blocked_end_user(): + """ + Regression for a blocked customer being admitted by core auth: /customer/block writes + LiteLLM_EndUserTable.blocked=True, so common_checks must reject the request even when the + optional enterprise blocked-user callback isn't configured. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + with pytest.raises(Exception) as exc_info: + await common_checks( + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=None, + end_user_object=LiteLLM_EndUserTable(user_id="blocked-demo", blocked=True), + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=UserAPIKeyAuth(token="test-token"), + request=MagicMock(spec=Request), + ) + + assert "End user=blocked-demo is blocked" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_common_checks_allows_unblocked_end_user(): + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + result = await common_checks( + request_body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=None, + end_user_object=LiteLLM_EndUserTable(user_id="allowed-demo", blocked=False), + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=UserAPIKeyAuth(token="test-token"), + request=MagicMock(spec=Request), + ) + + 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 0af5ad6cd9b..ad3a8d19a6c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -6,6 +6,7 @@ from fastapi.responses import JSONResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient +import litellm from litellm.proxy._types import ( LiteLLM_EndUserTable, LitellmUserRoles, @@ -781,3 +782,107 @@ def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth): "deleted_customers": 2, "message": "Successfully deleted customers with ids: ['c1', 'c2']", } + + +@pytest.fixture +def mock_end_user_cache(): + with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock: + mock.async_delete_cache = AsyncMock() + yield mock + + +def _deleted_cache_keys(mock_cache: MagicMock) -> set: + return {call.kwargs["key"] for call in mock_cache.async_delete_cache.call_args_list} + + +def test_block_customer_invalidates_cached_end_user_objects( + mock_prisma_client, mock_user_api_key_auth, mock_end_user_cache +): + """ + Regression: a warmed cache entry kept a just-blocked customer usable. Core auth caches the row + under `end_user_id:{id}` and the enterprise blocked-user hook under `litellm:end_user_id:{id}`, + so /customer/block has to drop both. + """ + mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock( + return_value=LiteLLM_EndUserTable(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 _deleted_cache_keys(mock_end_user_cache) == { + "end_user_id:blocked-1", + "litellm:end_user_id:blocked-1", + } + + +def test_unblock_customer_persists_blocked_false_and_invalidates_cache( + mock_prisma_client, mock_user_api_key_auth, mock_end_user_cache, monkeypatch +): + """ + Regression: /customer/unblock only mutated the in-memory litellm.blocked_user_list, so the + durable blocked=True written by /customer/block survived and re-blocked the customer once the + enterprise cache expired. + """ + monkeypatch.setattr(litellm, "blocked_user_list", ["blocked-1", "other"]) + mock_prisma_client.db.litellm_endusertable.update_many = AsyncMock(return_value=1) + + response = client.post( + "/customer/unblock", + json={"user_ids": ["blocked-1"]}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200 + assert response.json() == {"blocked_users": ["other"]} + mock_prisma_client.db.litellm_endusertable.update_many.assert_awaited_once_with( + where={"user_id": {"in": ["blocked-1"]}}, + data={"blocked": False}, + ) + assert _deleted_cache_keys(mock_end_user_cache) == { + "end_user_id:blocked-1", + "litellm:end_user_id:blocked-1", + } + + +def test_unblock_customer_without_enterprise_callback_still_persists( + mock_prisma_client, mock_user_api_key_auth, mock_end_user_cache, monkeypatch +): + """ + /customer/block works without the enterprise callback, so its inverse must too; previously this + returned 400 and left blocked=True in the database. + """ + monkeypatch.setattr(litellm, "blocked_user_list", None) + mock_prisma_client.db.litellm_endusertable.update_many = AsyncMock(return_value=1) + + response = client.post( + "/customer/unblock", + json={"user_ids": ["blocked-1"]}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200 + assert response.json() == {"blocked_users": []} + mock_prisma_client.db.litellm_endusertable.update_many.assert_awaited_once() + + +def test_unblock_customer_rejects_filepath_blocked_user_list( + mock_prisma_client, mock_user_api_key_auth, mock_end_user_cache, monkeypatch +): + """A file-backed blocked_user_list can't be edited at runtime, so nothing is written at all.""" + monkeypatch.setattr(litellm, "blocked_user_list", "/etc/litellm/blocked_users.txt") + mock_prisma_client.db.litellm_endusertable.update_many = AsyncMock(return_value=1) + + response = client.post( + "/customer/unblock", + json={"user_ids": ["blocked-1"]}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 400 + mock_prisma_client.db.litellm_endusertable.update_many.assert_not_awaited() + assert _deleted_cache_keys(mock_end_user_cache) == set()