From 2bbc941cbdb1e16795932fd678d19ae8dfa2ae61 Mon Sep 17 00:00:00 2001 From: lewis Date: Sat, 10 Oct 2026 08:28:35 -0400 Subject: [PATCH 01/56] fix(proxy): persist customer block and unblock state --- .../enterprise_hooks/blocked_user_list.py | 7 +- .../proxy/common_utils/user_api_key_cache.py | 4 + .../customer_endpoints.py | 79 ++-- .../test_customer_endpoints.py | 388 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 5 files changed, 432 insertions(+), 50 deletions(-) diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py index f6998d58cd5..25a027b899c 100644 --- a/enterprise/enterprise_hooks/blocked_user_list.py +++ b/enterprise/enterprise_hooks/blocked_user_list.py @@ -7,7 +7,7 @@ ## This accepts a list of user id's for whom calls will be rejected -from typing import Literal, Optional +from typing import Final, Literal, Optional from fastapi import HTTPException @@ -17,12 +17,13 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, end_user_block_cache_key from litellm.proxy.utils import PrismaClient class ENTERPRISE_BlockedUserList(CustomLogger): enforces_request_content: bool = True + blocked_user_list: list[str] | None # Class variables or attributes def __init__(self, prisma_client: Optional[PrismaClient]): self.prisma_client = prisma_client @@ -88,7 +89,7 @@ class ENTERPRISE_BlockedUserList(CustomLogger): }, ) - cache_key = f"litellm:end_user_id:{user}" + cache_key: Final = end_user_block_cache_key(f"{user}") end_user_cache_obj: Optional[LiteLLM_EndUserTable] = cache.get_cache( # type: ignore key=cache_key ) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 8c830816207..76ad9e9d17b 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -372,6 +372,10 @@ def end_user_cache_key(end_user_id: str) -> str: return f"end_user_id:{end_user_id}" +def end_user_block_cache_key(end_user_id: str) -> str: + return f"litellm:end_user_id:{end_user_id}" + + def end_user_restricted_registry_cache_key() -> str: """Cache key for the set of end-user ids whose row carries a restriction auth enforces.""" return "end_user_restricted_registry" diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index d604c58d1f5..a86c792ff6e 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_block_cache_key, end_user_cache_key, end_user_restricted_registry_cache_key, ) @@ -58,6 +59,7 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import ( _RowT_co: Final = TypeVar("_RowT_co", covariant=True) _STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) +_STRING_IDS: Final = TypeAdapter(tuple[str, ...]) _CLEARABLE_LIST_FIELDS: Final = frozenset({"models"}) @@ -105,6 +107,8 @@ if TYPE_CHECKING: data: Mapping[str, Mapping[str, object]], ) -> _RowT_co: ... + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + async def delete_many(self, where: Mapping[str, object]) -> int: ... @@ -135,7 +139,11 @@ async def _evict_end_user_cache_keys(cache_keys: Sequence[str]) -> None: def _end_user_cache_keys(user_ids: Sequence[str]) -> tuple[str, ...]: """The per-id entries plus the registry, which any restriction change can move ids in or out of.""" - return (*(end_user_cache_key(user_id) for user_id in user_ids), end_user_restricted_registry_cache_key()) + return ( + *(end_user_cache_key(user_id) for user_id in user_ids), + end_user_restricted_registry_cache_key(), + *(end_user_block_cache_key(user_id) for user_id in user_ids), + ) def _to_customer_response(record: BaseModel) -> CustomerResponse: @@ -160,7 +168,7 @@ def _to_customer_response(record: BaseModel) -> CustomerResponse: dependencies=[Depends(user_api_key_auth)], response_model=BlockUsersResponse, ) -async def block_user(data: BlockUsers): +async def block_user(data: BlockUsers) -> BlockUsersResponse: """ [BETA] Reject calls with this end-user id @@ -180,25 +188,26 @@ 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 _typed_table(EndUserRepository(prisma_client)).upsert( - where={"user_id": id}, - data={ - "create": {"user_id": id, "blocked": True}, - "update": {"blocked": True}, - }, - ) - records.append(record) - await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) - else: + if prisma_client is None: raise HTTPException( status_code=500, detail={"error": "Postgres DB Not connected"}, ) - - return {"blocked_users": records} + table: Final = _typed_table(EndUserRepository(prisma_client)) + records: Final = [ + _to_customer_response( + await table.upsert( + where={"user_id": user_id}, + data={ + "create": {"user_id": user_id, "blocked": True}, + "update": {"blocked": True}, + }, + ) + ) + for user_id in data.user_ids + ] + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + return BlockUsersResponse.model_validate({"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)}) @@ -216,7 +225,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 @@ -229,10 +238,10 @@ async def unblock_user(data: BlockUsers): }' ``` """ + from litellm.proxy.proxy_server import prisma_client + try: - from enterprise.enterprise_hooks.blocked_user_list import ( - ENTERPRISE_BlockedUserList, - ) + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList except ImportError: raise HTTPException( status_code=400, @@ -241,26 +250,34 @@ async def unblock_user(data: BlockUsers): + 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 - ): + configured_hook: Final = next( + (callback for callback in litellm.callbacks if isinstance(callback, ENTERPRISE_BlockedUserList)), None + ) + if configured_hook is None: raise HTTPException( status_code=400, detail={"error": "Blocked user check was never set. This call has no effect."}, ) + requested_ids: Final = frozenset(data.user_ids) + configured_ids: Final = _STRING_IDS.validate_python( + configured_hook.blocked_user_list + if isinstance(litellm.blocked_user_list, str) + else litellm.blocked_user_list or () + ) + remaining_legacy_ids: Final = tuple(user_id for user_id in configured_ids if user_id not in requested_ids) + table: Final = _typed_table(EndUserRepository(prisma_client)) if prisma_client is not None else None + if table is not None: + await table.update_many(where={"user_id": {"in": data.user_ids}}, data={"blocked": False}) if isinstance(litellm.blocked_user_list, list): - for id in data.user_ids: - litellm.blocked_user_list.remove(id) - else: + litellm.blocked_user_list[:] = remaining_legacy_ids + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + if isinstance(litellm.blocked_user_list, str) and requested_ids.intersection(configured_ids): raise HTTPException( status_code=500, detail={"error": "`blocked_user_list` must be set as a list. Filepaths can't be updated."}, ) - - return {"blocked_users": litellm.blocked_user_list} + return UnblockUsersResponse(blocked_users=list(remaining_legacy_ids)) def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index c99db2d8054..ee8d404c8f5 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1,4 +1,7 @@ from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6,14 +9,22 @@ from fastapi import FastAPI, HTTPException, Request, status from fastapi.responses import JSONResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient +from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow from litellm.proxy._types import ( + BlockUsers, LiteLLM_EndUserTable, LitellmUserRoles, ProxyException, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.management_endpoints.customer_endpoints import _should_update_field, router +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + end_user_block_cache_key, + end_user_cache_key, + end_user_restricted_registry_cache_key, +) +from litellm.proxy.management_endpoints.customer_endpoints import _should_update_field, block_user, router, unblock_user from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -546,25 +557,27 @@ def test_update_customer_response_keeps_nested_budget_server_fields(mock_prisma_ assert "updated_by" not in budget -def test_block_customer_success_serializes_through_response_model(mock_prisma_client, mock_user_api_key_auth): +def test_block_customer_success_serializes_through_response_model( + mock_prisma_client: MagicMock, mock_user_api_key_auth: object +) -> None: """ /customer/block returns {"blocked_users": []}. With response_model=BlockUsersResponse, a shape mismatch would raise a 500 ResponseValidationError, so a clean 200 proves the model matches runtime output. """ - blocked_row = LiteLLM_EndUserTable(user_id="blocked-1", blocked=True) + blocked_row: Final = PrismaEndUserRow(user_id="blocked-1", blocked=True, spend=0, models=[]) mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(return_value=blocked_row) - response = client.post( + response: Final = client.post( "/customer/block", json={"user_ids": ["blocked-1"]}, headers={"Authorization": "Bearer test-key"}, ) assert response.status_code == 200 - body = response.json() - assert body["blocked_users"][0]["user_id"] == "blocked-1" - assert body["blocked_users"][0]["blocked"] is True + body: Final = BlockUsersResponse.model_validate(response.json()) + assert body.blocked_users[0].user_id == blocked_row.user_id + assert body.blocked_users[0].blocked is True def test_delete_customer_success_serializes_through_response_model(mock_prisma_client, mock_user_api_key_auth): @@ -915,9 +928,7 @@ def test_customer_new_forwards_models_to_db(mock_prisma_client, mock_user_api_ke @pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) -def test_customer_new_rejects_a_duration_that_never_advances( - mock_prisma_client, mock_user_api_key_auth, bad_duration -): +def test_customer_new_rejects_a_duration_that_never_advances(mock_prisma_client, mock_user_api_key_auth, bad_duration): """A zero-length window resets to "now", leaving the customer's budget row permanently due for the reset job to re-read every tick.""" mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) @@ -1073,8 +1084,8 @@ def test_customer_new_invalidates_end_user_and_registry_caches(mock_prisma_clien ) assert response.status_code == 200, response.text - assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry"] - assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry"] + assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry", "litellm:end_user_id:c1"] + assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry", "litellm:end_user_id:c1"] def test_customer_update_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): @@ -1092,8 +1103,8 @@ def test_customer_update_invalidates_end_user_and_registry_caches(mock_prisma_cl ) assert response.status_code == 200, response.text - assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry"] - assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry"] + assert recording_cache.deleted == ["end_user_id:c1", "end_user_restricted_registry", "litellm:end_user_id:c1"] + assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry", "litellm:end_user_id:c1"] def test_customer_block_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): @@ -1114,14 +1125,359 @@ def test_customer_block_invalidates_end_user_and_registry_caches(mock_prisma_cli "end_user_id:c1", "end_user_id:c2", "end_user_restricted_registry", + "litellm:end_user_id:c1", + "litellm:end_user_id:c2", ] assert _published_keys(mock_publish) == [ "end_user_id:c1", "end_user_id:c2", "end_user_restricted_registry", + "litellm:end_user_id:c1", + "litellm:end_user_id:c2", ] +def test_customer_block_legacy_route_matches_canonical_response( + mock_prisma_client: MagicMock, mock_user_api_key_auth: object +) -> None: + blocked_row: Final = PrismaEndUserRow( + user_id="legacy-customer", blocked=True, alias="Legacy customer", spend=0, models=[] + ) + upsert: Final = AsyncMock(return_value=blocked_row) + mock_prisma_client.db = SimpleNamespace(litellm_endusertable=SimpleNamespace(upsert=upsert)) + cache: Final = UserApiKeyCache() + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + canonical: Final = client.post( + "/customer/block", json={"user_ids": [blocked_row.user_id]}, headers={"Authorization": "Bearer k"} + ) + legacy: Final = client.post( + "/end_user/block", json={"user_ids": [blocked_row.user_id]}, headers={"Authorization": "Bearer k"} + ) + assert canonical.status_code == legacy.status_code == 200 + assert legacy.json() == canonical.json() + body: Final = BlockUsersResponse.model_validate(legacy.json()) + assert len(body.blocked_users) == 1 + assert body.blocked_users[0].user_id == blocked_row.user_id + assert body.blocked_users[0].blocked is True + assert upsert.await_count == 2 + + +@pytest.mark.parametrize("legacy_ids", (None, (), ("legacy-customer", "still-blocked"))) +def test_customer_unblock_legacy_route_matches_canonical_response( + mock_prisma_client: MagicMock, mock_user_api_key_auth: object, legacy_ids: tuple[str, ...] | None +) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + update_many: Final = AsyncMock(return_value=1) + mock_prisma_client.db = SimpleNamespace(litellm_endusertable=SimpleNamespace(update_many=update_many)) + configured_ids: Final = list(legacy_ids) if legacy_ids is not None else None + cache: Final = UserApiKeyCache() + with ( + patch("litellm.blocked_user_list", configured_ids), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + with patch("litellm.callbacks", [hook]): + canonical: Final = client.post( + "/customer/unblock", json={"user_ids": ["legacy-customer"]}, headers={"Authorization": "Bearer k"} + ) + legacy: Final = client.post( + "/end_user/unblock", json={"user_ids": ["legacy-customer"]}, headers={"Authorization": "Bearer k"} + ) + assert canonical.status_code == legacy.status_code == 200 + assert legacy.json() == canonical.json() + body: Final = UnblockUsersResponse.model_validate(legacy.json()) + assert body.blocked_users == (["still-blocked"] if legacy_ids else []) + update_many.assert_awaited_with(where={"user_id": {"in": ["legacy-customer"]}}, data={"blocked": False}) + assert update_many.await_count == 2 + + +@pytest.mark.parametrize("legacy_ids", ((), ("c1", "legacy-only", "c2"))) +def test_customer_unblock_persists_database_and_evicts_both_auth_caches( + mock_prisma_client: MagicMock, + mock_user_api_key_auth: object, + legacy_ids: tuple[str, ...], +) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + update_many: Final = AsyncMock(return_value=2) + find_many: Final = AsyncMock(side_effect=AssertionError("Unblock must not read unrelated database customers")) + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace(update_many=update_many, find_many=find_many) + ) + configured_ids: Final = list(legacy_ids) + with ( + patch("litellm.blocked_user_list", configured_ids), + _end_user_cache_doubles() as (recording_cache, mock_publish), + ): + legacy_hook: Final = ENTERPRISE_BlockedUserList(prisma_client=None) + with patch("litellm.callbacks", [legacy_hook]): + response: Final = client.post( + "/customer/unblock", + json={"user_ids": ["c1", "c2"]}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + update_many.assert_awaited_once_with(where={"user_id": {"in": ["c1", "c2"]}}, data={"blocked": False}) + body: Final = UnblockUsersResponse.model_validate(response.json()) + assert body.blocked_users == (["legacy-only"] if legacy_ids else []) + assert legacy_hook.blocked_user_list == (["legacy-only"] if legacy_ids else []) + expected_keys: Final = ( + "end_user_id:c1", + "end_user_id:c2", + "end_user_restricted_registry", + "litellm:end_user_id:c1", + "litellm:end_user_id:c2", + ) + assert tuple(recording_cache.deleted) == expected_keys + assert tuple(_published_keys(mock_publish)) == expected_keys + + +@pytest.mark.asyncio +async def test_customer_block_unblock_refreshes_enterprise_hook_cache(mock_prisma_client: MagicMock) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + active_row: Final = PrismaEndUserRow(user_id="cached-customer", blocked=False, spend=0, models=[]) + blocked_row: Final = active_row.model_copy(update={"blocked": True}) + read_customer: Final = AsyncMock(side_effect=(active_row, blocked_row, active_row)) + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace( + find_unique=read_customer, + upsert=AsyncMock(return_value=blocked_row), + update_many=AsyncMock(return_value=1), + ) + ) + cache: Final = UserApiKeyCache() + customer_ids: Final = BlockUsers(user_ids=[active_row.user_id]) + auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.blocked_user_list", []), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + await hook.async_pre_call_hook(auth, cache, {"user": active_row.user_id}, "completion") + cached_row: Final = cache.get_cache(end_user_block_cache_key(active_row.user_id), model_type=PrismaEndUserRow) + assert cached_row is not None and cached_row.blocked is False + + await block_user(customer_ids) + with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): + await hook.async_pre_call_hook(auth, cache, {"user": active_row.user_id}, "completion") + + with patch("litellm.callbacks", [hook]): + await unblock_user(customer_ids) + await hook.async_pre_call_hook(auth, cache, {"user": active_row.user_id}, "completion") + + assert read_customer.await_count == 3 + + +@pytest.mark.parametrize( + "operation,initially_blocked,finally_blocked", + ( + ("new", False, True), + ("update", False, True), + ("update", True, False), + ("delete", True, False), + ), +) +@pytest.mark.asyncio +async def test_customer_mutations_refresh_warm_enterprise_hook_cache( + mock_prisma_client: MagicMock, + mock_user_api_key_auth: object, + operation: str, + initially_blocked: bool, + finally_blocked: bool, +) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + before_row: Final = PrismaEndUserRow(user_id="mutated-customer", blocked=initially_blocked, spend=0, models=[]) + after_row: Final = before_row.model_copy(update={"blocked": finally_blocked}) + read_customer: Final = AsyncMock( + side_effect=(None if operation == "new" else before_row, None if operation == "delete" else after_row) + ) + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace( + find_unique=read_customer, + find_first=AsyncMock(return_value=before_row), + find_many=AsyncMock(return_value=[before_row]), + create=AsyncMock(return_value=after_row), + update=AsyncMock(return_value=after_row), + delete_many=AsyncMock(return_value=1), + ) + ) + cache: Final = UserApiKeyCache() + auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.blocked_user_list", []), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + if initially_blocked: + with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): + await hook.async_pre_call_hook(auth, cache, {"user": before_row.user_id}, "completion") + else: + await hook.async_pre_call_hook(auth, cache, {"user": before_row.user_id}, "completion") + cached_before: Final = cache.get_cache( + end_user_block_cache_key(before_row.user_id), model_type=PrismaEndUserRow + ) + assert cached_before is not None and cached_before.blocked is initially_blocked + + payload: Final = ( + {"user_ids": [before_row.user_id]} + if operation == "delete" + else {"user_id": before_row.user_id, "blocked": finally_blocked} + ) + response: Final = client.post(f"/customer/{operation}", json=payload, headers={"Authorization": "Bearer k"}) + assert response.status_code == 200, response.text + if finally_blocked: + with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): + await hook.async_pre_call_hook(auth, cache, {"user": before_row.user_id}, "completion") + else: + await hook.async_pre_call_hook(auth, cache, {"user": before_row.user_id}, "completion") + assert read_customer.await_count == 2 + + +@pytest.mark.parametrize("configured_hook,configured_ids", ((False, ()), (False, ("c1",)), (False, None), (True, None))) +def test_customer_unblock_requires_loaded_hook( + mock_prisma_client: MagicMock, + mock_user_api_key_auth: object, + configured_hook: bool, + configured_ids: tuple[str, ...] | None, +) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + update_many: Final = AsyncMock(return_value=1) + mock_prisma_client.db = SimpleNamespace(litellm_endusertable=SimpleNamespace(update_many=update_many)) + blocked_row: Final = PrismaEndUserRow(user_id="c1", blocked=True, spend=0, models=[]) + cache: Final = UserApiKeyCache() + cache.set_cache(end_user_block_cache_key(blocked_row.user_id), blocked_row) + cache.set_cache(end_user_cache_key(blocked_row.user_id), blocked_row) + cache.set_cache(end_user_restricted_registry_cache_key(), [blocked_row.user_id]) + with ( + patch("litellm.blocked_user_list", list(configured_ids) if configured_ids is not None else None), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + with patch("litellm.callbacks", [hook] if configured_hook else []): + response: Final = client.post( + "/customer/unblock", + json={"user_ids": [blocked_row.user_id]}, + headers={"Authorization": "Bearer k"}, + ) + + cached_row: Final = cache.get_cache(end_user_block_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow) + if configured_hook: + assert response.status_code == 200, response.text + assert response.json() == {"blocked_users": []} + update_many.assert_awaited_once_with(where={"user_id": {"in": [blocked_row.user_id]}}, data={"blocked": False}) + assert cached_row is None + assert cache.get_cache(end_user_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow) is None + assert cache.get_cache(end_user_restricted_registry_cache_key()) is None + else: + assert response.status_code == 400 + assert response.json() == {"detail": {"error": "Blocked user check was never set. This call has no effect."}} + update_many.assert_not_awaited() + assert cached_row is not None and cached_row.blocked is True + cached_customer: Final = cache.get_cache(end_user_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow) + assert cached_customer is not None and cached_customer.blocked is True + assert cache.get_cache(end_user_restricted_registry_cache_key()) == [blocked_row.user_id] + + +@pytest.mark.asyncio +async def test_customer_unblock_database_only_hook_refreshes_cached_block(mock_prisma_client: MagicMock) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + blocked_row: Final = PrismaEndUserRow(user_id="database-only-customer", blocked=True, spend=0, models=[]) + active_row: Final = blocked_row.model_copy(update={"blocked": False}) + read_customer: Final = AsyncMock(side_effect=(blocked_row, active_row)) + update_many: Final = AsyncMock(return_value=1) + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace(find_unique=read_customer, update_many=update_many) + ) + cache: Final = UserApiKeyCache() + auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + cache.set_cache(end_user_cache_key(blocked_row.user_id), blocked_row) + cache.set_cache(end_user_restricted_registry_cache_key(), [blocked_row.user_id]) + with ( + patch("litellm.blocked_user_list", None), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): + await hook.async_pre_call_hook(auth, cache, {"user": blocked_row.user_id}, "completion") + cached_before: Final = cache.get_cache( + end_user_block_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow + ) + assert cached_before is not None and cached_before.blocked is True + with patch("litellm.callbacks", [hook]): + response: Final = await unblock_user(BlockUsers(user_ids=[blocked_row.user_id])) + assert response.blocked_users == [] + update_many.assert_awaited_once_with(where={"user_id": {"in": [blocked_row.user_id]}}, data={"blocked": False}) + assert cache.get_cache(end_user_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow) is None + assert cache.get_cache(end_user_restricted_registry_cache_key()) is None + assert cache.get_cache(end_user_block_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow) is None + await hook.async_pre_call_hook(auth, cache, {"user": blocked_row.user_id}, "completion") + cached_after: Final = cache.get_cache( + end_user_block_cache_key(blocked_row.user_id), model_type=PrismaEndUserRow + ) + assert cached_after is not None and cached_after.blocked is False + assert read_customer.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("in_file", (False, True)) +async def test_customer_unblock_file_list_clears_database_and_keeps_file( + mock_prisma_client: MagicMock, tmp_path: Path, in_file: bool +) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + user_id: Final = "file-customer" + file_contents: Final = f"{user_id}\nlegacy-only" if in_file else "legacy-only" + block_file: Final = tmp_path / "blocked-users.txt" + block_file.write_text(file_contents) + blocked_row: Final = PrismaEndUserRow(user_id=user_id, blocked=True, spend=0, models=[]) + active_row: Final = blocked_row.model_copy(update={"blocked": False}) + update_many: Final = AsyncMock(return_value=1) + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace(update_many=update_many, find_unique=AsyncMock(return_value=active_row)) + ) + cache: Final = UserApiKeyCache() + cache.set_cache(end_user_block_cache_key(user_id), blocked_row) + auth: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.blocked_user_list", str(block_file)), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.redis_usage_cache", None), + ): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + with patch("litellm.callbacks", [hook]): + if in_file: + with pytest.raises(HTTPException, match="Filepaths can't be updated") as error: + await unblock_user(BlockUsers(user_ids=[user_id])) + assert error.value.status_code == 500 + else: + response: Final = await unblock_user(BlockUsers(user_ids=[user_id])) + assert response.blocked_users == ["legacy-only"] + + assert cache.get_cache(end_user_block_cache_key(user_id), model_type=PrismaEndUserRow) is None + if in_file: + with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): + await hook.async_pre_call_hook(auth, cache, {"user": user_id}, "completion") + else: + await hook.async_pre_call_hook(auth, cache, {"user": user_id}, "completion") + + update_many.assert_awaited_once_with(where={"user_id": {"in": [user_id]}}, data={"blocked": False}) + assert block_file.read_text() == file_contents + + def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): """Without this a deleted customer keeps its cached budget and block enforced until the TTL expires.""" mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( @@ -1144,9 +1500,13 @@ def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_cl "end_user_id:c1", "end_user_id:c2", "end_user_restricted_registry", + "litellm:end_user_id:c1", + "litellm:end_user_id:c2", ] assert _published_keys(mock_publish) == [ "end_user_id:c1", "end_user_id:c2", "end_user_restricted_registry", + "litellm:end_user_id:c1", + "litellm:end_user_id:c2", ] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 53ecf9b2d75..f1c9bc6a8b3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -56686,7 +56686,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["BlockUsersResponse"]; }; }; /** @description Validation Error */ @@ -57071,7 +57071,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["UnblockUsersResponse"]; }; }; /** @description Validation Error */ From e2b8e3261ef3b5eb3d037a6496a28952f648f281 Mon Sep 17 00:00:00 2001 From: lewis Date: Fri, 9 Oct 2026 10:01:06 -0400 Subject: [PATCH 02/56] fix(proxy): preserve unsent object permission fields --- .../object_permission_utils.py | 3 +- .../test_customer_endpoints.py | 39 +++++++++++ .../test_object_permission_utils.py | 70 ++++++++++++++++--- 3 files changed, 102 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 86545e6e064..969ebf2cf8c 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -113,9 +113,10 @@ async def prepare_object_permission_upsert( existing_mcp_tool_permissions=existing_fields.get("mcp_tool_permissions"), prisma_client=prisma_client, ) + # Omitted request fields become None; preserve stored grants and Prisma column defaults. merged: Final[dict[str, object]] = { **existing_fields, - **new_object_permission, + **{field: value for field, value in new_object_permission.items() if value is not None}, "object_permission_id": object_permission_id, } record: Final[dict[str, object]] = { diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index c99db2d8054..2ac7b68766b 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1,4 +1,6 @@ +from collections.abc import Mapping from contextlib import contextmanager +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -9,6 +11,7 @@ from fastapi.testclient import TestClient from litellm.proxy._types import ( LiteLLM_EndUserTable, + LiteLLM_ObjectPermissionTable, LitellmUserRoles, ProxyException, ) @@ -1006,6 +1009,42 @@ def test_customer_update_applies_nonempty_object_permission(mock_prisma_client, ) +def test_customer_update_creates_partial_toolset_permission( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None +) -> None: + async def persist_permission( + *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> LiteLLM_ObjectPermissionTable: + assert data["create"] == { + "object_permission_id": where["object_permission_id"], + "mcp_toolsets": ["toolset-id"], + } + assert data["update"] == data["create"] + return LiteLLM_ObjectPermissionTable.model_validate(data["create"]) + + mock_prisma_client.configure_mock( + **{ + "db.litellm_endusertable.find_first": AsyncMock( + return_value=LiteLLM_EndUserTable(user_id="customer-id", blocked=False) + ), + "db.litellm_objectpermissiontable.find_unique": AsyncMock(return_value=None), + "db.litellm_objectpermissiontable.upsert": AsyncMock(side_effect=persist_permission), + "db.litellm_endusertable.update": AsyncMock( + return_value=LiteLLM_EndUserTable(user_id="customer-id", blocked=False) + ), + } + ) + + response: Final = client.post( + "/customer/update", + json={"user_id": "customer-id", "object_permission": {"mcp_toolsets": ["toolset-id"]}}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, response.text + assert response.json()["user_id"] == "customer-id" + + def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth): mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( return_value=[ diff --git a/tests/unit/proxy/management_helpers/test_object_permission_utils.py b/tests/unit/proxy/management_helpers/test_object_permission_utils.py index 33557713162..34201bde6e0 100644 --- a/tests/unit/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/unit/proxy/management_helpers/test_object_permission_utils.py @@ -1519,25 +1519,77 @@ _SHARED_ALIAS_DB_SERVERS = ( ) -def _make_ambiguity_prisma(existing_tool_permissions=None): +def _make_ambiguity_prisma( + existing_tool_permissions: dict[str, list[str]] | None = None, + existing_row_override: LiteLLM_ObjectPermissionTable | None = None, +) -> MagicMock: """Mock prisma client whose MCP server table holds _SHARED_ALIAS_DB_SERVERS and whose object permission row (if any) stores the given mcp_tool_permissions JSON string.""" - mock_prisma = MagicMock() + mock_prisma: Final = MagicMock() mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=list(_SHARED_ALIAS_DB_SERVERS)) mock_prisma.db.litellm_objectpermissiontable.create = AsyncMock( return_value=MagicMock(object_permission_id="perm-id") ) - existing_row = None - if existing_tool_permissions is not None: - existing_row = MagicMock() - existing_row.model_dump.return_value = { - "object_permission_id": "perm-id", - "mcp_tool_permissions": json.dumps(existing_tool_permissions), - } + stored_row: Final = ( + MagicMock( + model_dump=MagicMock( + return_value={ + "object_permission_id": "perm-id", + "mcp_tool_permissions": json.dumps(existing_tool_permissions), + } + ) + ) + if existing_tool_permissions is not None + else None + ) + existing_row: Final = existing_row_override if existing_row_override is not None else stored_row mock_prisma.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_row) return mock_prisma +@pytest.mark.asyncio +async def test_prepare_object_permission_upsert_omits_unsent_fields_on_create(): + permission: Final = LiteLLM_ObjectPermissionBase(mcp_toolsets=["toolset-id"]) + mock_prisma: Final = _make_ambiguity_prisma() + + upsert: Final = await prepare_object_permission_upsert( + new_object_permission=permission.model_dump(), + existing_object_permission_id=None, + prisma_client=mock_prisma, + ) + + assert upsert.object_permission_id + assert upsert.record == { + "object_permission_id": upsert.object_permission_id, + "mcp_toolsets": permission.mcp_toolsets, + } + + +@pytest.mark.asyncio +async def test_prepare_object_permission_upsert_preserves_unsent_grants_on_update(): + existing_row: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-id", + mcp_servers=["solo-id"], + mcp_tool_permissions={"solo-id": ["tool1"]}, + ) + permission: Final = LiteLLM_ObjectPermissionBase(mcp_toolsets=["toolset-id"]) + mock_prisma: Final = _make_ambiguity_prisma(existing_row_override=existing_row) + + upsert: Final = await prepare_object_permission_upsert( + new_object_permission=permission.model_dump(), + existing_object_permission_id=existing_row.object_permission_id, + prisma_client=mock_prisma, + ) + + assert upsert.object_permission_id == existing_row.object_permission_id + assert upsert.record == { + "object_permission_id": existing_row.object_permission_id, + "mcp_servers": existing_row.mcp_servers, + "mcp_tool_permissions": json.dumps(existing_row.mcp_tool_permissions), + "mcp_toolsets": permission.mcp_toolsets, + } + + @pytest.mark.asyncio @pytest.mark.parametrize( "identifier, colliding_ids", From f2e41384ef8a13eb74cdfd4bb7fa2d01c1147b8d Mon Sep 17 00:00:00 2001 From: lewis Date: Sat, 10 Oct 2026 08:24:28 -0400 Subject: [PATCH 03/56] feat(ui): add customers management page --- .../_components/CreateCustomerModal.tsx | 42 ++ .../customers/_components/CustomerForm.tsx | 260 +++++++++ .../_components/CustomerInfoView.tsx | 272 ++++++++++ .../CustomersPanel.integration.test.tsx | 502 ++++++++++++++++++ .../customers/_components/CustomersPanel.tsx | 205 +++++++ .../customers/_components/CustomersTable.tsx | 85 +++ .../_components/CustomersTableColumns.test.ts | 56 ++ .../_components/CustomersTableColumns.tsx | 240 +++++++++ .../_components/customerPayload.test.ts | 302 +++++++++++ .../customers/_components/customerPayload.ts | 146 +++++ .../customers/_components/useSaveCustomer.ts | 17 + .../src/app/(dashboard)/customers/page.tsx | 7 + .../hooks/customers/customerApi.test.ts | 44 ++ .../hooks/customers/customerApi.ts | 26 + .../customers/useCustomerMutations.test.tsx | 126 +++++ .../hooks/customers/useCustomerMutations.ts | 36 ++ .../src/app/(dashboard)/legacyPageRoutes.ts | 1 + .../src/components/leftnav.tsx | 2 + .../MCPServerSelector.test.tsx | 14 + .../MCPServerSelector.tsx | 3 + .../MCPToolPermissions.test.tsx | 30 ++ .../MCPToolPermissions.tsx | 3 + .../components/object_permissions_view.tsx | 2 +- .../permissions/MCPServerPermissions.test.tsx | 1 + .../permissions/MCPServerPermissions.tsx | 4 +- 25 files changed, 2422 insertions(+), 4 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerModal.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerModal.tsx new file mode 100644 index 00000000000..877b326e602 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CreateCustomerModal.tsx @@ -0,0 +1,42 @@ +"use client"; + +import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { toast } from "@/lib/toast"; +import { CustomerForm } from "./CustomerForm"; +import type { CustomerFormValues } from "./customerPayload"; +import { useSaveCustomer } from "./useSaveCustomer"; + +type CreateCustomerModalProps = { + open: boolean; + onOpenChange: (open: boolean) => void; + onCreated?: (customerId: string) => void; +}; + +function CreateCustomerFields({ onOpenChange, onCreated }: Omit) { + const { saveCustomer, isSaving } = useSaveCustomer(null); + const handleSubmit = async (values: CustomerFormValues) => { + try { + const customer = await saveCustomer(values); + toast.success("Customer created successfully"); + onOpenChange(false); + onCreated?.(customer.user_id); + } catch (error) { + toast.fromError(error); + } + }; + + return onOpenChange(false)} isSaving={isSaving} />; +} + +export function CreateCustomerModal({ open, onOpenChange, onCreated }: CreateCustomerModalProps) { + return ( + + + + Create Customer + + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx new file mode 100644 index 00000000000..f10b2f06e45 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx @@ -0,0 +1,260 @@ +"use client"; + +import { z } from "zod"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useBudgetOptions } from "@/app/(dashboard)/hooks/budgets/useBudgetOptions"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { useUserModels } from "@/app/(dashboard)/hooks/models/useModels"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; +import { FormField } from "@/components/shared/form/FormField"; +import { MultiSelect } from "@/components/shared/MultiSelect"; +import { SearchSelect } from "@/components/shared/SearchSelect"; +import { Button } from "@/components/ui/button"; +import { FieldGroup } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { toCustomerFormValues, type CustomerFormValues } from "./customerPayload"; + +const mcpSelectionShape = { + servers: z.array(z.string()), + accessGroups: z.array(z.string()), + toolsets: z.array(z.string()), +}; + +const customerShape = { + user_id: z.string().trim().min(1, "Customer ID is required"), + alias: z.string(), + budget_id: z.string(), + max_budget: z + .string() + .refine( + (value) => !value.trim() || (Number.isFinite(Number(value)) && Number(value) >= 0), + "Enter a non-negative budget", + ), + models: z.array(z.string()), + allowed_model_region: z.enum(["", "eu", "us"]), + default_model: z.string(), + mcp_servers_and_groups: z.object(mcpSelectionShape), + mcp_tool_permissions: z.record(z.string(), z.array(z.string())), +}; + +const customerSchema = z.object(customerShape).refine((values) => !values.budget_id || !values.max_budget.trim(), { + message: "Select a budget or enter a max budget, not both", + path: ["max_budget"], +}); + +const editCustomerSchema = customerSchema.refine( + (values) => !values.max_budget.trim() || Number(values.max_budget) > 0, + { message: "Enter a budget greater than zero when editing", path: ["max_budget"] }, +); + +const regionOptions = [ + { value: "", label: "Any" }, + { value: "eu", label: "EU" }, + { value: "us", label: "US" }, +]; + +type CustomerFormProps = { + customer?: EndUser | null; + isEdit?: boolean; + isSaving?: boolean; + onSubmit: (values: CustomerFormValues) => void | Promise; + onCancel: () => void; +}; + +export function CustomerForm({ + customer = null, + isEdit = false, + isSaving = false, + onSubmit, + onCancel, +}: CustomerFormProps) { + const { accessToken } = useAuthorized(); + const { data: budgets = [] } = useBudgetOptions(accessToken); + const { data: models = [] } = useUserModels(); + const initialValues = toCustomerFormValues(customer); + const form = useZodForm(isEdit ? editCustomerSchema : customerSchema, { + defaultValues: initialValues, + }); + const hasLinkedBudget = isEdit && Boolean(initialValues.budget_id); + const budgetDescription = isEdit + ? "Leaving both budget fields empty keeps the current budget" + : "Choose a budget or enter a maximum below"; + const handleSubmit = form.handleSubmit(async (values) => { + if (hasLinkedBudget && !values.budget_id && values.max_budget.trim()) { + form.setError("max_budget", { + message: "This customer has a linked budget. Leave Max Budget empty or choose another budget", + }); + return; + } + await onSubmit(values); + }); + const selectedBudgetId = form.watch("budget_id"); + const budgetOptions = [ + { value: "", label: "No budget" }, + ...(selectedBudgetId && !budgets.some((budget) => budget.budget_id === selectedBudgetId) + ? [{ value: selectedBudgetId, label: selectedBudgetId }] + : []), + ...budgets.map((budget) => ({ value: budget.budget_id, label: budget.budget_id })), + ]; + const modelOptions = models.map((model) => ({ label: model, value: model })); + const defaultModelEmptyLabel = isEdit && initialValues.default_model ? "Keep current default" : "None"; + const defaultModelOptions = [{ value: "", label: defaultModelEmptyLabel }, ...modelOptions]; + const selection = form.watch("mcp_servers_and_groups"); + + return ( +
+ + + {({ ref, ...control }) => } + + + + {({ ref, ...control }) => } + + + + {({ id, value, onChange }) => ( + + )} + + + + {({ ref, ...control }) => ( + + )} + + + + {({ id, value, onChange }) => ( + + )} + + + + {({ id, value, onChange }) => ( + + )} + + + + {({ id, value, onChange }) => ( + onChange(selected ?? "")} + placeholder={defaultModelEmptyLabel} + emptyText="No models found" + /> + )} + + + + {({ id, value, onChange }) => ( + + )} + + form.setValue("mcp_tool_permissions", permissions)} + /> + +
+ + +
+
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx new file mode 100644 index 00000000000..a6da69ef4fc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx @@ -0,0 +1,272 @@ +"use client"; + +import { ArrowLeft, MoreVertical, Pencil, Trash2 } from "lucide-react"; +import type { ReactNode } from "react"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import ObjectPermissionsView from "@/components/object_permissions_view"; +import { getBudgetDurationLabel } from "@/components/common_components/budget_duration_dropdown"; +import CopyButton from "@/components/shared/CopyButton"; +import { PageHeader, PageHeaderTitle } from "@/components/shared/PageHeader"; +import { ModelsCell, MoneyCell } from "@/components/shared/table_cells"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { toast } from "@/lib/toast"; + +import { CustomerForm } from "./CustomerForm"; +import { CustomerBlockAction, CustomerBlockedBadge } from "./CustomersTableColumns"; +import type { CustomerFormValues } from "./customerPayload"; +import { useSaveCustomer } from "./useSaveCustomer"; + +type CustomerTab = "overview" | "settings"; + +type CustomerInfoViewProps = { + customer: EndUser; + accessToken: string | null; + canModify: boolean; + premiumUser: boolean; + isEdit: boolean; + tab: CustomerTab; + isBlocking: boolean; + onTabChange: (tab: CustomerTab) => void; + onBack: () => void; + onEdit: () => void; + onCancelEdit: () => void; + onDelete: (customer: EndUser) => void; + onSetBlocked: (customer: EndUser) => void; +}; + +function OverviewField({ label, children }: { label: string; children: ReactNode }) { + return ( +
+
{label}
+
{children}
+
+ ); +} + +function CustomerModelFields({ customer }: { customer: EndUser }) { + return ( + <> + + + + + {customer.allowed_model_region?.toUpperCase() ?? "Any"} + + {customer.default_model ?? "None"} + + ); +} + +function CustomerOverview({ customer, accessToken }: Pick) { + const budget = customer.litellm_budget_table; + return ( +
+
+ + + Budget Status + + +
+ +
+
+ of{" "} + + + +
+

Reset: {getBudgetDurationLabel(budget?.budget_duration)}

+ + + + + Rate Limits + + +
+ {budget?.tpm_limit ?? "Unlimited"} + {budget?.rpm_limit ?? "Unlimited"} +
+ + + + + Models + + +
+ +
+ + +
+ +
+ ); +} + +function CustomerSettings({ customer, accessToken }: Pick) { + const budget = customer.litellm_budget_table; + return ( +
+
+ {customer.alias || "Not set"} + {customer.budget_id ?? "No budget"} + + + + + + {getBudgetDurationLabel(budget?.budget_duration)} + {budget?.tpm_limit ?? "Unlimited"} + {budget?.rpm_limit ?? "Unlimited"} + +
+ +
+ ); +} + +function CustomerActions({ + customer, + premiumUser, + isBlocking, + onDelete, + onSetBlocked, +}: Pick) { + return ( + + }> + + + + + onDelete(customer)}> + + Delete Customer + + + + ); +} + +export default function CustomerInfoView({ + customer, + accessToken, + canModify, + isEdit, + tab, + onTabChange, + onBack, + onEdit, + onCancelEdit, + ...actions +}: CustomerInfoViewProps) { + const { saveCustomer, isSaving } = useSaveCustomer(customer); + const handleSave = async (values: CustomerFormValues) => { + try { + await saveCustomer(values); + toast.success("Customer updated"); + onCancelEdit(); + } catch (error) { + toast.fromError(error); + } + }; + const showEdit = canModify && isEdit; + const hasDistinctAlias = Boolean(customer.alias && customer.alias !== customer.user_id); + + return ( + <> + + +
+
+
+ + {customer.alias || customer.user_id} + + + {customer.blocked && } +
+ {hasDistinctAlias && ( +
+ Customer ID: {customer.user_id} + +
+ )} +
+ {canModify && !showEdit && } +
+ + { + if (value === "overview" || value === "settings") onTabChange(value); + }} + > + + + Overview + + + Settings + + + + + + +
+
+

Customer Settings

+ {canModify && !showEdit && ( + + )} +
+ {showEdit ? ( + + ) : ( + + )} +
+ + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx new file mode 100644 index 00000000000..31df2c883f6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx @@ -0,0 +1,502 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { Toaster, toast as sonner } from "sonner"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { ApiError } from "@/lib/http/client"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { chooseSelectOption, renderWithProviders, testQueryClient } from "../../../../../tests/test-utils"; +import CustomersPanel from "./CustomersPanel"; + +const { post, get, useQuery, fetchServers, listTools, refetch, authorization } = vi.hoisted(() => ({ + post: vi.fn(), + get: vi.fn(), + useQuery: vi.fn(), + fetchServers: vi.fn(), + listTools: vi.fn(), + refetch: vi.fn(), + authorization: { accessToken: "sk-test", userRole: "Admin", userId: "admin", isViewOnly: false, premiumUser: false }, +})); + +vi.unmock("@/lib/toast"); + +vi.mock("@/lib/http/api", () => ({ $api: { useQuery } })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => authorization })); +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + apiClient: { post, get }, + fetchMCPServers: fetchServers, + listMCPTools: listTools, + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), +})); + +const activeCustomer: EndUser = { + user_id: "alice@example.test", + alias: "Alice", + blocked: false, + spend: 12.5, + models: [], + object_permission: { + object_permission_id: "permission-alice", + mcp_servers: [], + mcp_access_groups: ["research"], + mcp_toolsets: ["toolset-1"], + mcp_tool_permissions: {}, + agents: [], + agent_access_groups: [], + models: [], + search_tools: [], + vector_stores: [], + blocked_tools: [], + }, +}; + +const blockedCustomer: EndUser = { + user_id: "bob@example.test", + alias: "Bob", + blocked: true, + spend: 0, + models: [], +}; + +const customers = [activeCustomer, blockedCustomer]; + +describe("Customers management", () => { + beforeEach(() => { + vi.clearAllMocks(); + sonner.dismiss(); + testQueryClient.clear(); + authorization.userRole = "Admin"; + authorization.isViewOnly = false; + authorization.premiumUser = false; + get.mockResolvedValue([]); + post.mockResolvedValue(activeCustomer); + fetchServers.mockResolvedValue([]); + listTools.mockResolvedValue({ tools: [], error: false }); + useQuery.mockImplementation( + (_method: string, path: string, init: { params?: { query?: { end_user_id?: string } } }) => ({ + data: + path === "/customer/list" + ? customers + : customers.find((customer) => customer.user_id === init.params?.query?.end_user_id), + isLoading: false, + isError: false, + isFetching: false, + refetch, + }), + ); + }); + + it("refreshes the customer list", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Refresh" })); + expect(refetch).toHaveBeenCalledOnce(); + }); + + it("uses the shared money display for zero and small spend", () => { + useQuery.mockReturnValue({ + data: [{ ...activeCustomer, spend: 0.001 }, blockedCustomer], + isLoading: false, + isError: false, + }); + renderWithProviders(); + expect(within(screen.getByRole("row", { name: /alice@example.test/ })).getByText("< $0.01")).toBeInTheDocument(); + expect( + within(screen.getByRole("row", { name: /bob@example.test/ })).queryByText(/\$0\.00/), + ).not.toBeInTheDocument(); + }); + + it("keeps MCP access optional while showing each customer status", async () => { + const user = userEvent.setup(); + renderWithProviders(); + expect(screen.queryByRole("columnheader", { name: "MCP Access" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Columns" })); + await user.click(await screen.findByRole("menuitemcheckbox", { name: "MCP Access" })); + await user.keyboard("{Escape}"); + + const alice = screen.getByRole("row", { name: /alice@example.test/ }); + expect(within(alice).getByText("Active")).toBeInTheDocument(); + expect(within(alice).getByText("1 group, 1 toolset")).toBeInTheDocument(); + const bob = screen.getByRole("row", { name: /bob@example.test/ }); + expect(within(bob).getByText("Blocked")).toBeInTheDocument(); + expect(within(bob).getByText("No restriction")).toBeInTheDocument(); + }); + + it.each(["row", "detail"])("disables Block without Enterprise in the %s menu", async (surface) => { + const user = userEvent.setup(); + renderWithProviders( + , + surface === "detail" ? { searchParams: "customer=alice%40example.test" } : undefined, + ); + await user.click( + screen.getByRole("button", { + name: surface === "detail" ? "More customer actions" : "Actions for alice@example.test", + }), + ); + const block = await screen.findByRole("menuitem", { name: "Block" }); + expect(block).toHaveAttribute("aria-disabled", "true"); + await user.hover(screen.getByLabelText("Block requires Enterprise")); + expect(await screen.findByRole("tooltip")).toHaveTextContent( + "This is a LiteLLM Enterprise feature, and requires a valid key to use.", + ); + await user.click(screen.getByLabelText("Block requires Enterprise")); + expect(post).not.toHaveBeenCalled(); + }); + + it.each(["row", "detail"])("blocks through customer update with Enterprise in the %s menu", async (surface) => { + authorization.premiumUser = true; + const user = userEvent.setup(); + renderWithProviders( + , + surface === "detail" ? { searchParams: "customer=alice%40example.test" } : undefined, + ); + await user.click( + screen.getByRole("button", { + name: surface === "detail" ? "More customer actions" : "Actions for alice@example.test", + }), + ); + const block = await screen.findByRole("menuitem", { name: "Block" }); + expect(block).not.toHaveAttribute("aria-disabled", "true"); + await user.click(block); + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: activeCustomer.user_id, blocked: true }, + }), + ); + }); + + it.each([false, true])("unblocks from the header with premiumUser=%s", async (premiumUser) => { + authorization.premiumUser = premiumUser; + const user = userEvent.setup(); + renderWithProviders(, { searchParams: "customer=bob%40example.test" }); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Unblock" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "More customer actions" })); + const unblock = await screen.findByRole("menuitem", { name: "Unblock" }); + expect(unblock).not.toHaveAttribute("aria-disabled", "true"); + expect(screen.getByRole("menuitem", { name: "Delete Customer" })).toBeInTheDocument(); + await user.click(unblock); + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: blockedCustomer.user_id, blocked: false }, + }), + ); + }); + + it.each([null, activeCustomer.user_id])("shows the customer ID once when alias is %s", (alias) => { + useQuery.mockReturnValue({ data: { ...activeCustomer, alias }, isLoading: false, isError: false }); + renderWithProviders(, { searchParams: "customer=alice%40example.test" }); + expect(screen.getAllByText(activeCustomer.user_id)).toHaveLength(1); + expect(screen.getByRole("heading", { name: /alice@example.test/ })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Copy Customer ID" })).toBeInTheDocument(); + }); + + it.each(["Alice", "alice@example.test"])("filters by ID or alias using %s", async (search) => { + renderWithProviders(); + fireEvent.change(screen.getByRole("textbox", { name: "Search customers" }), { target: { value: search } }); + + await waitFor(() => expect(screen.queryByRole("row", { name: /bob@example.test/ })).not.toBeInTheDocument()); + expect(screen.getByRole("row", { name: /alice@example.test/ })).toBeInTheDocument(); + }); + + it("creates a customer with the filled form values", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Create Customer" })); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByLabelText("Customer ID"), { target: { value: " new@example.test " } }); + fireEvent.change(within(dialog).getByLabelText("Alias"), { target: { value: "New Customer" } }); + fireEvent.change(within(dialog).getByLabelText("Max Budget (USD)"), { target: { value: "25" } }); + await user.click(within(dialog).getByRole("button", { name: "Create Customer" })); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/new", { + accessToken: "sk-test", + body: { user_id: "new@example.test", alias: "New Customer", max_budget: 25 }, + }), + ); + }); + + it("requires a customer ID before creating", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Create Customer" })); + const dialog = await screen.findByRole("dialog"); + await user.click(within(dialog).getByRole("button", { name: "Create Customer" })); + + expect(await within(dialog).findByText("Customer ID is required")).toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + }); + + it("edits a deep-linked customer without rewriting MCP permissions", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: "customer=alice%40example.test", onUrlUpdate }); + expect(screen.getByRole("heading", { name: /^Alice/ })).toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: "Edit Settings" })); + const alias = await screen.findByLabelText("Alias"); + expect(screen.getByLabelText("Customer ID")).toBeDisabled(); + fireEvent.change(alias, { target: { value: "Updated Alice" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: activeCustomer.user_id, alias: "Updated Alice", models: [] }, + }), + ); + await waitFor(() => + expect(onUrlUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + searchParams: expect.any(URLSearchParams), + }), + ), + ); + expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("edit")).toBeNull(); + }); + + it("unblocks a customer from the row menu", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("button", { name: "Actions for bob@example.test" })); + await user.click(await screen.findByRole("menuitem", { name: "Unblock" })); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: blockedCustomer.user_id, blocked: false }, + }), + ); + }); + + it("preserves an unrestricted server when editing only the alias after tools load", async () => { + const user = userEvent.setup(); + const directCustomer: EndUser = { + ...activeCustomer, + object_permission: { + ...activeCustomer.object_permission!, + mcp_servers: ["server-1"], + mcp_access_groups: [], + mcp_toolsets: [], + }, + }; + useQuery.mockReturnValue({ data: directCustomer, isLoading: false, isError: false }); + fetchServers.mockResolvedValue([ + { + server_id: "server-1", + server_name: "Research Server", + created_at: "", + created_by: "admin", + updated_at: "", + updated_by: "admin", + }, + ]); + listTools.mockResolvedValue({ + error: false, + tools: [ + { name: "read_document", description: "Read a document", inputSchema: {} }, + { name: "delete_document", description: "Delete a document", inputSchema: {} }, + ], + }); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + await screen.findByRole("radio", { name: "Flat List" }); + fireEvent.change(screen.getByLabelText("Alias"), { target: { value: "Updated Alice" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: activeCustomer.user_id, alias: "Updated Alice", models: [] }, + }), + ); + }); + + it("rejects an edit budget of zero because customer updates cannot apply it", async () => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "0" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect(await screen.findByText("Enter a budget greater than zero when editing")).toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + }); + + it("confirms deletion and clears the open customer from the URL", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: "customer=alice%40example.test", onUrlUpdate }); + await user.click(screen.getByRole("button", { name: "More customer actions" })); + await user.click(await screen.findByRole("menuitem", { name: "Delete Customer" })); + const dialog = await screen.findByRole("dialog", { name: "Delete Customer?" }); + expect(within(dialog).getByText(/recreates the customer with no restrictions/)).toBeInTheDocument(); + const confirm = within(dialog).getByRole("button", { name: "Delete" }); + expect(confirm).toBeDisabled(); + fireEvent.change(within(dialog).getByPlaceholderText(activeCustomer.user_id), { + target: { value: activeCustomer.user_id }, + }); + await user.click(confirm); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/delete", { + accessToken: "sk-test", + body: { user_ids: [activeCustomer.user_id] }, + }), + ); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalledWith(expect.objectContaining({ queryString: "" }))); + expect(screen.getByRole("heading", { name: "Customers" })).toBeInTheDocument(); + }); + + it("keeps Admin Viewer read-only, including an edit deep link", async () => { + authorization.isViewOnly = true; + const user = userEvent.setup(); + const list = renderWithProviders(); + expect(screen.queryByRole("button", { name: "Create Customer" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Actions for alice@example.test" })); + expect(await screen.findByRole("menuitem", { name: "Copy ID" })).toBeInTheDocument(); + expect(screen.queryByRole("menuitem", { name: "Edit" })).not.toBeInTheDocument(); + list.unmount(); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + await user.click(screen.getByRole("tab", { name: "Settings" })); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + }); + it("matches the Organizations money, model and rate-limit columns", () => { + const customer: EndUser = { + ...activeCustomer, + budget_id: "shared-budget", + litellm_budget_table: { + budget_id: "shared-budget", + max_budget: 100, + tpm_limit: 600, + rpm_limit: 6, + created_at: "2026-01-01T00:00:00Z", + }, + }; + useQuery.mockReturnValue({ data: [customer], isLoading: false, isError: false }); + renderWithProviders(); + + expect(screen.getByText("Manage end-user spend, budgets, rate limits, model and MCP access")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Spend (USD)" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Budget (USD)" })).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "TPM / RPM Limits" })).toBeInTheDocument(); + const row = screen.getByRole("row", { name: /alice@example.test/ }); + expect(within(row).getByText("$12.50")).toBeInTheDocument(); + expect(within(row).getByText("$100.00")).toBeInTheDocument(); + expect(within(row).getByText("TPM: 600")).toBeInTheDocument(); + expect(within(row).getByText("RPM: 6")).toBeInTheDocument(); + expect(within(row).getByText("All Proxy Models")).toBeInTheDocument(); + }); + + it("shows Overview cards and shared permissions before entering Settings", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + const customer: EndUser = { + ...activeCustomer, + allowed_model_region: "eu", + default_model: "default-model", + litellm_budget_table: { + max_budget: 100, + budget_duration: "30d", + tpm_limit: 600, + rpm_limit: 6, + created_at: "2026-01-01T00:00:00Z", + }, + }; + useQuery.mockReturnValue({ data: customer, isLoading: false, isError: false }); + renderWithProviders(, { searchParams: "customer=alice%40example.test", onUrlUpdate }); + + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByText("Budget Status")).toBeInTheDocument(); + expect(screen.getByText("$100.00")).toBeInTheDocument(); + expect(screen.getByText("Reset: monthly")).toBeInTheDocument(); + expect(screen.getByText("Rate Limits")).toBeInTheDocument(); + expect(screen.getByText("Object Permissions")).toBeInTheDocument(); + expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + expect(screen.getByText("EU")).toBeInTheDocument(); + expect(screen.getByText("default-model")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("tab", { name: "Settings" })); + expect(await screen.findByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("customer_tab")).toBe("settings"); + expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("customer")).toBe(activeCustomer.user_id); + expect(onUrlUpdate.mock.calls.at(-1)?.[0].options.history).toBe("push"); + }); + + it("opens a Settings deep link without enabling the edit form", () => { + renderWithProviders(, { searchParams: "customer=alice%40example.test&customer_tab=settings" }); + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByRole("button", { name: "Edit Settings" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); + }); + + it("rejects a max budget after clearing a shared budget link", async () => { + const user = userEvent.setup(); + const budget = { budget_id: "shared-budget", max_budget: 100, created_at: "2026-01-01T00:00:00Z" }; + const linkedCustomers: EndUser[] = customers.map((customer) => ({ + ...customer, + budget_id: budget.budget_id, + litellm_budget_table: budget, + })); + get.mockResolvedValue([budget]); + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? linkedCustomers : linkedCustomers[0], + isLoading: false, + isError: false, + })); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "true"); + await chooseSelectOption(user, screen.getByLabelText("Budget"), "No budget"); + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "25" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect( + await screen.findByText("This customer has a linked budget. Leave Max Budget empty or choose another budget"), + ).toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: activeCustomer.user_id, alias: "Alice", models: [] }, + }), + ); + }); + + it("shows a failed customer update as an error toast", async () => { + const user = userEvent.setup(); + const message = "Customer update failed"; + post.mockRejectedValue(new ApiError(message, 400, { detail: { error: message } })); + renderWithProviders( + <> + + + , + { searchParams: "customer=bob%40example.test" }, + ); + await user.click(screen.getByRole("button", { name: "More customer actions" })); + await user.click(await screen.findByRole("menuitem", { name: "Unblock" })); + + expect(await screen.findByText(message)).toBeInTheDocument(); + expect(screen.getByText("Request Error")).toBeInTheDocument(); + expect(screen.queryByText("Customer unblocked")).not.toBeInTheDocument(); + expect(screen.getByText("Blocked")).toBeInTheDocument(); + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: blockedCustomer.user_id, blocked: false }, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx new file mode 100644 index 00000000000..43306e508e2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx @@ -0,0 +1,205 @@ +"use client"; + +import { Contact, Plus, RefreshCw } from "lucide-react"; +import { parseAsBoolean, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; +import { useCallback, useState } from "react"; + +import { + useCustomer, + useDeleteCustomers, + useSetCustomerBlockedState, +} from "@/app/(dashboard)/hooks/customers/useCustomerMutations"; +import { useCustomers, type EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { Alert, AlertTitle } from "@/components/shared/Alert"; +import { Page } from "@/components/shared/Page"; +import { PageHeader, PageHeaderControls, PageHeaderDescription, PageHeaderTitle } from "@/components/shared/PageHeader"; +import { Button } from "@/components/ui/button"; +import { toast } from "@/lib/toast"; +import { all_admin_roles, isProxyAdminRole } from "@/utils/roles"; + +import { CreateCustomerModal } from "./CreateCustomerModal"; +import CustomerInfoView from "./CustomerInfoView"; +import CustomersTable from "./CustomersTable"; + +const CUSTOMER_VIEW_PARSERS = { + customer: parseAsString, + edit: parseAsBoolean.withDefault(false), + customer_tab: parseAsStringLiteral(["overview", "settings"]).withDefault("overview"), +}; + +function CustomersHeader({ + canModify, + onCreate, + onRefresh, + isRefreshing, +}: { + canModify: boolean; + onCreate: () => void; + onRefresh: () => void; + isRefreshing: boolean; +}) { + return ( + + + + Customers + + Manage end-user spend, budgets, rate limits, model and MCP access + + + {canModify && ( + + )} + + + ); +} + +export default function CustomersPanel() { + const { accessToken, userRole, isViewOnly, premiumUser } = useAuthorized(); + const canModify = isProxyAdminRole(userRole ?? "") && !isViewOnly; + const [{ customer: customerId, edit, customer_tab: tab }, setView] = useQueryStates(CUSTOMER_VIEW_PARSERS, { + history: "push", + }); + const [isCreateOpen, setIsCreateOpen] = useState(false); + const [deleteTarget, setDeleteTarget] = useState(null); + const customers = useCustomers(); + const customer = useCustomer(customerId); + const deleteCustomers = useDeleteCustomers(); + const setBlockedState = useSetCustomerBlockedState(); + + const openCustomer = useCallback( + (id: string) => void setView({ customer: id, edit: null, customer_tab: null }), + [setView], + ); + const editCustomer = useCallback( + (id: string) => void setView({ customer: id, edit: true, customer_tab: "settings" }), + [setView], + ); + const closeCustomer = useCallback(() => void setView({ customer: null, edit: null, customer_tab: null }), [setView]); + const cancelEdit = useCallback(() => void setView({ edit: null, customer_tab: "settings" }), [setView]); + const handleSetBlocked = useCallback( + async (target: EndUser) => { + if (!canModify || (!target.blocked && !premiumUser)) return; + try { + await setBlockedState.mutateAsync({ userId: target.user_id, blocked: !target.blocked }); + toast.success(target.blocked ? "Customer unblocked" : "Customer blocked"); + } catch (error) { + toast.fromError(error); + } + }, + [canModify, premiumUser, setBlockedState], + ); + const handleDelete = async () => { + if (!deleteTarget || !canModify) return; + try { + await deleteCustomers.mutateAsync([deleteTarget.user_id]); + if (customerId === deleteTarget.user_id) closeCustomer(); + setDeleteTarget(null); + toast.success("Customer deleted"); + } catch (error) { + toast.fromError(error); + } + }; + + if (!all_admin_roles.includes(userRole ?? "")) { + return ( + + + You need an admin role to view customers + + + ); + } + + return ( + + {customerId ? ( + void setView({ customer_tab: value, edit: null })} + isBlocking={setBlockedState.isPending} + onBack={closeCustomer} + onEdit={() => editCustomer(customerId)} + onCancelEdit={cancelEdit} + onDelete={setDeleteTarget} + onSetBlocked={handleSetBlocked} + /> + ) : ( + <> + setIsCreateOpen(true)} + onRefresh={() => void customers.refetch()} + isRefreshing={customers.isFetching} + /> + + + )} + {canModify && ( + <> + + setDeleteTarget(null)} + onOk={handleDelete} + confirmLoading={deleteCustomers.isPending} + /> + + )} + + ); +} + +type CustomerDetailProps = Omit[0], "customer"> & { + customer: EndUser | undefined; + isLoading: boolean; + isError: boolean; +}; + +function CustomerDetail({ customer, isLoading, isError, onBack, ...props }: CustomerDetailProps) { + if (isLoading) return

Loading customer

; + if (isError || !customer) { + return ( + <> + + + Unable to load customer + + + ); + } + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx new file mode 100644 index 00000000000..f0f5e71368b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx @@ -0,0 +1,85 @@ +"use client"; + +import { Contact } from "lucide-react"; +import { useMemo } from "react"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { DataTable, DataTableViewOptions, useUrlTableState } from "@/components/shared/DataTable"; +import { Input } from "@/components/ui/input"; + +import { getCustomerTableColumns, type CustomerTableActions } from "./CustomersTableColumns"; + +const TABLE_STATE_OPTIONS = { + sortFields: ["user_id", "spend", "max_budget", "blocked"], + defaultSort: { id: "user_id", desc: false }, + defaultPageSize: 50, + filterColumns: [], +}; + +const DEFAULT_COLUMN_VISIBILITY = { mcp_access: false }; + +type CustomersTableProps = CustomerTableActions & { + data: EndUser[]; + isLoading: boolean; + isError: boolean; +}; + +function EmptyCustomers() { + return ( +
+ +

No customers found

+

Create a customer to manage end-user limits and access

+
+ ); +} + +export default function CustomersTable({ data, isLoading, isError, ...actions }: CustomersTableProps) { + const { search, setSearch, sorting, onSortingChange, pagination, onPaginationChange } = + useUrlTableState(TABLE_STATE_OPTIONS); + const { canModify, premiumUser, onView, onEdit, onDelete, onSetBlocked } = actions; + const columns = useMemo(() => { + const columnActions = { canModify, premiumUser, onView, onEdit, onDelete, onSetBlocked }; + return getCustomerTableColumns(columnActions); + }, [canModify, premiumUser, onView, onEdit, onDelete, onSetBlocked]); + const filteredCustomers = useMemo(() => { + const query = search.trim().toLowerCase(); + return data.filter((customer) => `${customer.user_id} ${customer.alias ?? ""}`.toLowerCase().includes(query)); + }, [data, search]); + + return ( +
+ customer.user_id} + defaultColumnVisibility={DEFAULT_COLUMN_VISIBILITY} + toolbar={(table) => ( +
+ setSearch(event.target.value)} + /> + +
+ )} + paginationMode="client" + pagination={pagination} + onPaginationChange={onPaginationChange} + sortingMode="client" + sorting={sorting} + onSortingChange={onSortingChange} + isLoading={isLoading} + isError={isError} + loadingMessage="Loading customers" + noDataMessage={} + size="compact" + /> +
+ ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.ts new file mode 100644 index 00000000000..eec60502c9c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.ts @@ -0,0 +1,56 @@ +import { describe, expect, it } from "vitest"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; + +import { describeMcpGrant } from "./CustomersTableColumns"; + +const emptyGrant: NonNullable = { + object_permission_id: "permission-id", + agent_access_groups: [], + agents: [], + blocked_tools: [], + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: {}, + models: [], + search_tools: [], + vector_stores: [], +}; + +describe("describeMcpGrant", () => { + it.each([null, undefined, emptyGrant])("describes an empty grant as unrestricted", (permission) => { + expect(describeMcpGrant(permission)).toBe("No restriction"); + }); + + it("summarizes servers, groups, toolsets and tool limits", () => { + const permission = { + ...emptyGrant, + mcp_servers: ["server-a", "server-b"], + mcp_access_groups: ["group-id"], + mcp_toolsets: ["toolset-id"], + mcp_tool_permissions: { "server-a": ["read"] }, + }; + expect(describeMcpGrant(permission)).toBe("2 servers, 1 group, 1 toolset, 1 tool limit"); + }); + + it("does not describe a tool-only grant as unrestricted", () => { + const permission = { ...emptyGrant, mcp_tool_permissions: { "server-id": ["read"] } }; + expect(describeMcpGrant(permission)).toBe("1 tool limit"); + }); + + it("labels the all-servers sentinel instead of counting it as a server", () => { + const permission = { ...emptyGrant, mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL] }; + expect(describeMcpGrant(permission)).toBe("All MCP servers"); + }); + + it("shows the no-servers sentinel as blocked even if other grants exist", () => { + const permission = { + ...emptyGrant, + mcp_servers: [NO_MCP_SERVERS_SENTINEL], + mcp_access_groups: ["group-id"], + }; + expect(describeMcpGrant(permission)).toBe("No MCP access"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx new file mode 100644 index 00000000000..775273ad43a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx @@ -0,0 +1,240 @@ +"use client"; + +import type { ColumnDef } from "@tanstack/react-table"; +import { Ban, Copy, MoreHorizontal, Pencil, ShieldCheck, Trash2 } from "lucide-react"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { IdentityCell, ModelsCell, MoneyCell, StatusBadge } from "@/components/shared/table_cells"; +import { buttonVariants } from "@/components/ui/button"; +import { Badge } from "@/components/ui/badge"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/cva.config"; +import { copyToClipboard } from "@/utils/dataUtils"; + +export type CustomerTableActions = { + canModify: boolean; + premiumUser: boolean; + onView: (customerId: string) => void; + onEdit: (customerId: string) => void; + onDelete: (customer: EndUser) => void; + onSetBlocked: (customer: EndUser) => void; +}; + +export function CustomerBlockedBadge() { + return ( + + + }> + + + Blocked + + + + {"Enforced only when LiteLLM Enterprise's blocked user check is enabled"} + + + + ); +} + +export function CustomerBlockAction({ + customer, + premiumUser, + isBlocking = false, + onSetBlocked, +}: { + customer: EndUser; + premiumUser: boolean; + isBlocking?: boolean; + onSetBlocked: (customer: EndUser) => void; +}) { + const action = ( + onSetBlocked(customer)} + > + {customer.blocked ? : } + {customer.blocked ? "Unblock" : "Block"} + + ); + if (customer.blocked || premiumUser) return action; + return ( + + + }> + {action} + + + This is a LiteLLM Enterprise feature, and requires a valid key to use. + + + + ); +} + +const countLabel = (count: number, noun: string): string => (count ? `${count} ${noun}${count === 1 ? "" : "s"}` : ""); + +export function describeMcpGrant(permission: EndUser["object_permission"]): string { + const servers = permission?.mcp_servers ?? []; + if (servers.includes(NO_MCP_SERVERS_SENTINEL)) return "No MCP access"; + const serverLabel = servers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL) + ? "All MCP servers" + : countLabel(servers.length, "server"); + const parts = [ + serverLabel, + countLabel(permission?.mcp_access_groups?.length ?? 0, "group"), + countLabel(permission?.mcp_toolsets?.length ?? 0, "toolset"), + countLabel(Object.keys(permission?.mcp_tool_permissions ?? {}).length, "tool limit"), + ].filter(Boolean); + return parts.join(", ") || "No restriction"; +} + +function CustomerRowActions({ + customer, + canModify, + premiumUser, + onEdit, + onDelete, + onSetBlocked, +}: CustomerTableActions & { customer: EndUser }) { + return ( + + + + + + {canModify && ( + <> + onEdit(customer.user_id)}> + + Edit + + + + )} + void copyToClipboard(customer.user_id, "Customer ID copied")}> + + Copy ID + + {canModify && ( + <> + + onDelete(customer)}> + + Delete + + + )} + + + ); +} + +export function getCustomerTableColumns(actions: CustomerTableActions): ColumnDef[] { + return [ + { + id: "user_id", + accessorKey: "user_id", + meta: { title: "Customer ID", skeleton: "twoLine" }, + header: ({ column }) => , + size: 260, + cell: ({ row }) => ( + actions.onView(row.original.user_id)} + /> + ), + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend (USD)", numeric: true }, + header: ({ column }) => , + size: 130, + cell: ({ row }) => , + }, + { + id: "max_budget", + accessorFn: (customer) => customer.litellm_budget_table?.max_budget, + meta: { title: "Budget (USD)", numeric: true }, + header: ({ column }) => , + size: 140, + cell: ({ row }) => ( + + ), + }, + { + id: "models", + meta: { title: "Models", skeleton: "chips" }, + header: "Models", + size: 240, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "limits", + meta: { title: "TPM / RPM Limits" }, + header: "TPM / RPM Limits", + size: 150, + enableSorting: false, + cell: ({ row }) => + row.original.litellm_budget_table?.tpm_limit == null && row.original.litellm_budget_table?.rpm_limit == null ? ( + "-" + ) : ( +
+ TPM: {row.original.litellm_budget_table?.tpm_limit ?? "Unlimited"} + RPM: {row.original.litellm_budget_table?.rpm_limit ?? "Unlimited"} +
+ ), + }, + { + id: "mcp_access", + meta: { title: "MCP Access" }, + header: "MCP Access", + size: 210, + enableSorting: false, + cell: ({ row }) => {describeMcpGrant(row.original.object_permission)}, + }, + { + id: "blocked", + accessorKey: "blocked", + meta: { title: "Status", skeleton: "badge" }, + header: ({ column }) => , + size: 120, + cell: ({ row }) => + row.original.blocked ? : , + }, + { + id: "actions", + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, + ]; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts new file mode 100644 index 00000000000..33d68688e2f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts @@ -0,0 +1,302 @@ +import { describe, expect, it } from "vitest"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import { buildCustomerPayload, toCustomerFormValues } from "./customerPayload"; + +const customer: EndUser = { user_id: "customer-id", blocked: false, spend: 0 }; +const permission = { + object_permission_id: "permission-id", + mcp_servers: ["server-id"], + mcp_access_groups: ["group-id"], + mcp_toolsets: ["toolset-id"], + mcp_tool_permissions: { "server-id": ["read", "write"] }, + agent_access_groups: [], + agents: [], + blocked_tools: [], + models: [], + search_tools: [], + skills: [], + vector_stores: [], +}; + +const server: MCPServer = { + server_id: "server-id", + created_at: "2026-01-01T00:00:00Z", + created_by: "admin", + updated_at: "2026-01-01T00:00:00Z", + updated_by: "admin", +}; + +describe("customer form values", () => { + it("starts with empty optional values and grants", () => { + const expected = { + user_id: "", + alias: "", + budget_id: "", + max_budget: "", + models: [], + allowed_model_region: "", + default_model: "", + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: {}, + }; + expect(toCustomerFormValues(null)).toEqual(expected); + }); + + it("loads the stored identity, models and MCP grant", () => { + const stored: EndUser = { + ...customer, + alias: "Alias", + budget_id: "budget-id", + models: ["model-id"], + allowed_model_region: "eu", + default_model: "default-model", + object_permission: permission, + litellm_budget_table: { max_budget: 50, created_at: "2026-01-01T00:00:00Z" }, + }; + const expected = { + user_id: stored.user_id, + alias: stored.alias, + budget_id: stored.budget_id, + max_budget: "", + models: stored.models, + allowed_model_region: stored.allowed_model_region, + default_model: stored.default_model, + mcp_servers_and_groups: { + servers: permission.mcp_servers, + accessGroups: permission.mcp_access_groups, + toolsets: permission.mcp_toolsets, + }, + mcp_tool_permissions: permission.mcp_tool_permissions, + }; + expect(toCustomerFormValues(stored)).toEqual(expected); + }); +}); + +describe("customer payload", () => { + it("creates a customer with only the fields that are set", () => { + const values = { ...toCustomerFormValues(null), user_id: " customer-id " }; + expect(buildCustomerPayload(values, null, [], [])).toEqual({ user_id: customer.user_id }); + }); + + it("keeps zero budgets and trims optional text fields", () => { + const values = { + ...toCustomerFormValues(null), + user_id: customer.user_id, + alias: " Alias ", + max_budget: "0", + models: ["model-id"], + allowed_model_region: "us" as const, + default_model: " default-model ", + }; + const expected = { + user_id: customer.user_id, + alias: "Alias", + max_budget: 0, + models: values.models, + allowed_model_region: values.allowed_model_region, + default_model: "default-model", + }; + expect(buildCustomerPayload(values, null, [], [])).toEqual(expected); + }); + + it("sends a selected existing budget", () => { + const values = { ...toCustomerFormValues(null), user_id: customer.user_id, budget_id: "budget-id" }; + expect(buildCustomerPayload(values, null, [], [])).toEqual({ + user_id: customer.user_id, + budget_id: values.budget_id, + }); + }); + + it.each(["", " "])("clears a stored alias when the form contains %j", (alias) => { + const stored = { ...customer, alias: "Alias" }; + const values = { ...toCustomerFormValues(stored), alias }; + + expect(buildCustomerPayload(values, stored, [], [])).toEqual({ + user_id: stored.user_id, + alias: "", + models: [], + }); + }); + + it("includes the grant when a toolset is chosen on create", () => { + const values = { + ...toCustomerFormValues(null), + user_id: customer.user_id, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["toolset-id"] }, + }; + const grant = { + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: values.mcp_servers_and_groups.toolsets, + mcp_tool_permissions: {}, + }; + expect(buildCustomerPayload(values, null, [], [])).toEqual({ user_id: customer.user_id, object_permission: grant }); + }); + + it("omits the stored MCP grant on an unrelated edit, regardless of ordering", () => { + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: ["server-id", "another-server"], + mcp_access_groups: ["group-id", "another-group"], + mcp_toolsets: ["toolset-id", "another-toolset"], + mcp_tool_permissions: { "server-id": ["read", "write"], "another-server": ["list"] }, + }, + }; + const values = { + ...toCustomerFormValues(stored), + alias: "New alias", + mcp_servers_and_groups: { + servers: ["another-server", "server-id"], + accessGroups: ["another-group", "group-id"], + toolsets: ["another-toolset", "toolset-id"], + }, + mcp_tool_permissions: { "another-server": ["list"], "server-id": ["write", "read"] }, + }; + expect(buildCustomerPayload(values, stored, [server], [])).toEqual({ + user_id: customer.user_id, + alias: values.alias, + models: [], + }); + }); + + it("keeps stored tool restrictions untouched when only the alias changes", () => { + const stored = { + ...customer, + object_permission: { ...permission, mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [] }, + }; + const values = { ...toCustomerFormValues(stored), alias: "New alias" }; + expect(buildCustomerPayload(values, stored, [server], [])).toEqual({ + user_id: customer.user_id, + alias: values.alias, + models: [], + }); + }); + + it("saves edits to a server granted only through its tool permissions", () => { + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-id": ["read"] }, + }, + }; + const values = { ...toCustomerFormValues(stored), mcp_tool_permissions: { "server-id": ["write"] } }; + const grant = { + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: values.mcp_tool_permissions, + }; + + expect(buildCustomerPayload(values, stored, [server], [])).toEqual({ + user_id: customer.user_id, + models: [], + object_permission: grant, + }); + }); + + it.each([ + { kind: "server", allowedTools: ["read"] }, + { kind: "toolset", allowedTools: ["write"] }, + ])("keeps an independent tool grant when adding a $kind", ({ kind, allowedTools }) => { + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-id": ["read"] }, + }, + }; + const anotherServer = { ...server, server_id: "another-server" }; + const toolsets: MCPToolset[] = [ + { + toolset_id: "toolset-id", + toolset_name: "Toolset", + tools: [{ server_id: anotherServer.server_id, tool_name: "list" }], + }, + ]; + const selection = { + servers: kind === "server" ? [anotherServer.server_id] : [], + accessGroups: [], + toolsets: kind === "toolset" ? [toolsets[0].toolset_id] : [], + }; + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: selection, + mcp_tool_permissions: { "server-id": allowedTools }, + }; + const grant = { + mcp_servers: selection.servers, + mcp_access_groups: [], + mcp_toolsets: selection.toolsets, + mcp_tool_permissions: values.mcp_tool_permissions, + }; + + expect(buildCustomerPayload(values, stored, [server, anotherServer], toolsets)).toEqual({ + user_id: customer.user_id, + models: [], + object_permission: grant, + }); + }); + + it("keeps a canonical server edit from an independent shared-alias grant when adding a server", () => { + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { wiki: ["read"] }, + }, + }; + const sharedServer = { ...server, alias: "wiki" }; + const anotherSharedServer = { ...sharedServer, server_id: "another-shared-server" }; + const addedServer = { ...server, server_id: "added-server" }; + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [addedServer.server_id], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { wiki: ["read"], [sharedServer.server_id]: ["write"] }, + }; + const grant = { + mcp_servers: values.mcp_servers_and_groups.servers, + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: values.mcp_tool_permissions, + }; + + expect(buildCustomerPayload(values, stored, [sharedServer, anotherSharedServer, addedServer], [])).toEqual({ + user_id: customer.user_id, + models: [], + object_permission: grant, + }); + }); + + it("sends an empty models list when clearing the selection on update", () => { + const stored = { ...customer, models: ["model-id"] }; + const values = { ...toCustomerFormValues(stored), models: [] }; + expect(buildCustomerPayload(values, stored, [], [])).toEqual({ user_id: customer.user_id, models: [] }); + }); + + it("drops the tool restriction for a deselected server", () => { + const stored = { ...customer, object_permission: { ...permission, mcp_access_groups: [], mcp_toolsets: [] } }; + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + }; + const grant = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; + expect(buildCustomerPayload(values, stored, [server], [])).toEqual({ + user_id: customer.user_id, + models: [], + object_permission: grant, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts new file mode 100644 index 00000000000..de0a49d6eab --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts @@ -0,0 +1,146 @@ +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import type { UpdateCustomerRequest } from "@/app/(dashboard)/hooks/customers/customerApi"; +import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import { mcpServersForIdentifier } from "@/components/mcp_server_management/effectiveMcpServers"; +import { extractMcpEntitlement, type McpEntitlementUpdate } from "@/components/mcp_server_management/mcpEntitlement"; + +export type CustomerFormValues = { + user_id: string; + alias: string; + budget_id: string; + max_budget: string; + models: string[]; + allowed_model_region: "" | "eu" | "us"; + default_model: string; + mcp_servers_and_groups: { servers: string[]; accessGroups: string[]; toolsets: string[] }; + mcp_tool_permissions: Record; +}; + +const storedMcpGrant = (customer: EndUser | null): McpEntitlementUpdate => ({ + mcp_servers: customer?.object_permission?.mcp_servers ?? [], + mcp_access_groups: customer?.object_permission?.mcp_access_groups ?? [], + mcp_toolsets: customer?.object_permission?.mcp_toolsets ?? [], + mcp_tool_permissions: customer?.object_permission?.mcp_tool_permissions ?? {}, +}); + +export const toCustomerFormValues = (customer: EndUser | null): CustomerFormValues => { + const grant = storedMcpGrant(customer); + const maxBudget = customer?.litellm_budget_table?.max_budget; + return { + user_id: customer?.user_id ?? "", + alias: customer?.alias ?? "", + budget_id: customer?.budget_id ?? "", + max_budget: customer?.budget_id || maxBudget == null ? "" : String(maxBudget), + models: customer?.models ?? [], + allowed_model_region: customer?.allowed_model_region ?? "", + default_model: customer?.default_model ?? "", + mcp_servers_and_groups: { + servers: grant.mcp_servers, + accessGroups: grant.mcp_access_groups, + toolsets: grant.mcp_toolsets, + }, + mcp_tool_permissions: grant.mcp_tool_permissions, + }; +}; + +const normalizedStrings = (values: string[]): string[] => Array.from(new Set(values)).toSorted(); + +const grantIdentity = (grant: McpEntitlementUpdate): string => { + const normalized = { + mcp_servers: normalizedStrings(grant.mcp_servers), + mcp_access_groups: normalizedStrings(grant.mcp_access_groups), + mcp_toolsets: normalizedStrings(grant.mcp_toolsets), + mcp_tool_permissions: Object.fromEntries( + Object.keys(grant.mcp_tool_permissions) + .toSorted() + .map((server) => [server, normalizedStrings(grant.mcp_tool_permissions[server])]), + ), + }; + return JSON.stringify(normalized); +}; + +const selectionIdentity = (grant: McpEntitlementUpdate): string => + grantIdentity({ ...grant, mcp_tool_permissions: {} }); + +const extractEditedMcpGrant = ( + values: CustomerFormValues, + storedGrant: McpEntitlementUpdate, + servers: MCPServer[], + toolsets: MCPToolset[], +): McpEntitlementUpdate | null => { + const grant = extractMcpEntitlement(values, servers, toolsets); + if (!grant) return null; + const previousValues = { + mcp_servers_and_groups: { + servers: storedGrant.mcp_servers, + accessGroups: storedGrant.mcp_access_groups, + toolsets: storedGrant.mcp_toolsets, + }, + mcp_tool_permissions: storedGrant.mcp_tool_permissions, + }; + const previousPermissions = extractMcpEntitlement(previousValues, servers, toolsets)?.mcp_tool_permissions ?? {}; + const independentKeys = new Set( + Object.keys(storedGrant.mcp_tool_permissions).filter((key) => !Object.hasOwn(previousPermissions, key)), + ); + const independentServerIds = new Set( + Array.from(independentKeys).flatMap((key) => + mcpServersForIdentifier(servers, key).map((server) => server.server_id), + ), + ); + const independentPermissions = Object.fromEntries( + Object.entries(values.mcp_tool_permissions).filter( + ([key]) => + independentKeys.has(key) || + mcpServersForIdentifier(servers, key).some((server) => independentServerIds.has(server.server_id)), + ), + ); + return { ...grant, mcp_tool_permissions: { ...grant.mcp_tool_permissions, ...independentPermissions } }; +}; + +const changedMcpGrant = ( + values: CustomerFormValues, + existing: EndUser | null, + servers: MCPServer[], + toolsets: MCPToolset[], +): McpEntitlementUpdate | null => { + const selection = values.mcp_servers_and_groups; + const formGrant = { + mcp_servers: selection.servers, + mcp_access_groups: selection.accessGroups, + mcp_toolsets: selection.toolsets, + mcp_tool_permissions: values.mcp_tool_permissions, + }; + const storedGrant = storedMcpGrant(existing); + const storedIdentity = grantIdentity(storedGrant); + const grant = + selectionIdentity(formGrant) === selectionIdentity(storedGrant) + ? formGrant + : extractEditedMcpGrant(values, storedGrant, servers, toolsets); + return grant && grantIdentity(grant) !== storedIdentity ? grant : null; +}; + +const optionalCustomerFields = ( + values: CustomerFormValues, + existing: EndUser | null, +): Partial => ({ + ...(values.alias.trim() || existing?.alias ? { alias: values.alias.trim() } : {}), + ...(values.budget_id ? { budget_id: values.budget_id } : {}), + ...(values.max_budget.trim() ? { max_budget: Number(values.max_budget) } : {}), + ...(values.allowed_model_region ? { allowed_model_region: values.allowed_model_region } : {}), + ...(values.default_model.trim() ? { default_model: values.default_model.trim() } : {}), +}); + +export const buildCustomerPayload = ( + values: CustomerFormValues, + existing: EndUser | null, + servers: MCPServer[], + toolsets: MCPToolset[], +): UpdateCustomerRequest => { + const grant = changedMcpGrant(values, existing, servers, toolsets); + return { + user_id: values.user_id.trim(), + ...optionalCustomerFields(values, existing), + ...(existing || values.models.length > 0 ? { models: values.models } : {}), + ...(grant ? { object_permission: grant } : {}), + }; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts new file mode 100644 index 00000000000..a0046ad8d25 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts @@ -0,0 +1,17 @@ +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { useCreateCustomer, useUpdateCustomer } from "@/app/(dashboard)/hooks/customers/useCustomerMutations"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { buildCustomerPayload, type CustomerFormValues } from "./customerPayload"; + +export const useSaveCustomer = (existing: EndUser | null) => { + const { data: servers = [] } = useMCPServers(); + const { data: toolsets = [] } = useMCPToolsets(); + const create = useCreateCustomer(); + const update = useUpdateCustomer(); + const saveCustomer = (values: CustomerFormValues) => { + const payload = buildCustomerPayload(values, existing, servers, toolsets); + return existing ? update.mutateAsync(payload) : create.mutateAsync(payload); + }; + return { saveCustomer, isSaving: create.isPending || update.isPending }; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx new file mode 100644 index 00000000000..1e5e37f65d8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/page.tsx @@ -0,0 +1,7 @@ +"use client"; + +import CustomersPanel from "./_components/CustomersPanel"; + +export default function CustomersPage() { + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts new file mode 100644 index 00000000000..057818a0ed3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts @@ -0,0 +1,44 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { apiClient } from "@/components/networking"; +import { createCustomer, deleteCustomers, setCustomerBlockedState, updateCustomer } from "./customerApi"; + +vi.mock("@/components/networking", () => ({ apiClient: { post: vi.fn() } })); + +const post = vi.mocked(apiClient.post); + +describe("customer API", () => { + beforeEach(() => post.mockReset()); + + it.each([ + { name: "create", action: createCustomer, path: "/customer/new" }, + { name: "update", action: updateCustomer, path: "/customer/update" }, + ])("posts the $name payload and returns the customer", async ({ action, path }) => { + const body = { user_id: "customer-id", alias: "Alias" }; + const response = { ...body, blocked: false, spend: 0 }; + post.mockResolvedValue(response); + + await expect(action("access-token", body)).resolves.toEqual(response); + expect(post).toHaveBeenCalledWith(path, { accessToken: "access-token", body }); + }); + + it("posts the customer IDs to delete", async () => { + const userIds = ["customer-id", "another-id"]; + const response = { deleted_customers: userIds.length, message: "Customers deleted" }; + post.mockResolvedValue(response); + + await expect(deleteCustomers("access-token", userIds)).resolves.toEqual(response); + expect(post).toHaveBeenCalledWith("/customer/delete", { accessToken: "access-token", body: { user_ids: userIds } }); + }); + + it.each([true, false])("updates the customer's blocked flag to %s", async (blocked) => { + const userId = "customer-id"; + const response = { user_id: userId, blocked, spend: 0 }; + post.mockResolvedValue(response); + + await expect(setCustomerBlockedState("access-token", { userId, blocked })).resolves.toEqual(response); + expect(post).toHaveBeenCalledWith("/customer/update", { + accessToken: "access-token", + body: { user_id: userId, blocked }, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts new file mode 100644 index 00000000000..731c9b2dfa6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts @@ -0,0 +1,26 @@ +import { apiClient } from "@/components/networking"; +import type { components } from "@/lib/http/schema"; +import type { EndUser } from "./useCustomers"; + +type Schemas = components["schemas"]; + +export type CreateCustomerRequest = Partial & { user_id: string }; +export type UpdateCustomerRequest = Partial & { user_id: string }; +export type CustomerBlockedStateInput = { userId: string; blocked: boolean }; + +export const createCustomer = (accessToken: string, body: CreateCustomerRequest): Promise => + apiClient.post("/customer/new", { accessToken, body }); + +export const updateCustomer = (accessToken: string, body: UpdateCustomerRequest): Promise => + apiClient.post("/customer/update", { accessToken, body }); + +export const deleteCustomers = (accessToken: string, userIds: string[]): Promise => + apiClient.post("/customer/delete", { + accessToken, + body: { user_ids: userIds }, + }); + +export const setCustomerBlockedState = ( + accessToken: string, + { userId, blocked }: CustomerBlockedStateInput, +): Promise => updateCustomer(accessToken, { user_id: userId, blocked }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx new file mode 100644 index 00000000000..f6500fd088a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx @@ -0,0 +1,126 @@ +import React, { type ReactNode } from "react"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { apiClient } from "@/components/networking"; +import { $api } from "@/lib/http/api"; +import { + useCreateCustomer, + useCustomer, + useDeleteCustomers, + useSetCustomerBlockedState, + useUpdateCustomer, +} from "./useCustomerMutations"; + +const authorized = vi.hoisted(() => vi.fn()); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: authorized })); +vi.mock("@/components/networking", () => ({ apiClient: { post: vi.fn() } })); +vi.mock("@/lib/http/api", () => ({ $api: { useQuery: vi.fn() } })); + +const post = vi.mocked(apiClient.post); +const query = vi.mocked($api.useQuery); + +const createWrapper = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false }, mutations: { retry: false } } }); + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ); + return { queryClient, wrapper }; +}; + +beforeEach(() => { + vi.clearAllMocks(); + authorized.mockReturnValue({ accessToken: "access-token", userRole: "Admin" }); +}); + +describe("useCustomer", () => { + it("fetches the requested customer with an admin access token", () => { + renderHook(() => useCustomer("customer-id")); + expect(query).toHaveBeenCalledWith( + "get", + "/customer/info", + { params: { query: { end_user_id: "customer-id" } } }, + { enabled: true }, + ); + }); + + it.each([ + { role: "Admin Viewer", accessToken: "access-token", id: "customer-id", enabled: true }, + { role: "Internal User", accessToken: "access-token", id: "customer-id", enabled: false }, + { role: "Admin", accessToken: null, id: "customer-id", enabled: false }, + { role: "Admin", accessToken: "access-token", id: null, enabled: false }, + ])("gates the query for $role with customer $id", ({ role, accessToken, id, enabled }) => { + authorized.mockReturnValue({ accessToken, userRole: role }); + renderHook(() => useCustomer(id)); + expect(query).toHaveBeenCalledWith( + "get", + "/customer/info", + { params: { query: { end_user_id: id ?? "" } } }, + { enabled }, + ); + }); +}); + +describe("customer mutations", () => { + it("invalidates both list and detail prefixes after a successful update", async () => { + const { queryClient, wrapper } = createWrapper(); + const listKey = ["get", "/customer/list", {}]; + const detailKey = ["get", "/customer/info", { params: { query: { end_user_id: "customer-id" } } }]; + const otherDetailKey = ["get", "/customer/info", { params: { query: { end_user_id: "another-id" } } }]; + const unrelatedKey = ["get", "/key/list", {}]; + queryClient.setQueryData(listKey, []); + queryClient.setQueryData(detailKey, {}); + queryClient.setQueryData(otherDetailKey, {}); + queryClient.setQueryData(unrelatedKey, []); + post.mockResolvedValue({ user_id: "customer-id" }); + const { result } = renderHook(() => useUpdateCustomer(), { wrapper }); + + await act(async () => { + await result.current.mutateAsync({ user_id: "customer-id", alias: "Alias" }); + }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(detailKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(otherDetailKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(unrelatedKey)?.isInvalidated).toBe(false); + }); + + it.each([ + { name: "create", useHook: useCreateCustomer }, + { name: "update", useHook: useUpdateCustomer }, + ])("rejects $name without an access token", async ({ useHook }) => { + authorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { wrapper } = createWrapper(); + const { result } = renderHook(useHook, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ user_id: "customer-id" })).rejects.toThrow("Access token is required"); + }); + expect(post).not.toHaveBeenCalled(); + }); + + it("rejects deletion without an access token", async () => { + authorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { wrapper } = createWrapper(); + const { result } = renderHook(useDeleteCustomers, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync(["customer-id"])).rejects.toThrow("Access token is required"); + }); + expect(post).not.toHaveBeenCalled(); + }); + + it("rejects block state changes without an access token", async () => { + authorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { wrapper } = createWrapper(); + const { result } = renderHook(useSetCustomerBlockedState, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ userId: "customer-id", blocked: true })).rejects.toThrow( + "Access token is required", + ); + }); + expect(post).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts new file mode 100644 index 00000000000..51c1cd8d9b0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts @@ -0,0 +1,36 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { $api } from "@/lib/http/api"; +import { all_admin_roles } from "@/utils/roles"; +import { createCustomer, deleteCustomers, setCustomerBlockedState, updateCustomer } from "./customerApi"; + +export const useCustomer = (id: string | null) => { + const { accessToken, userRole } = useAuthorized(); + return $api.useQuery( + "get", + "/customer/info", + { params: { query: { end_user_id: id ?? "" } } }, + { enabled: Boolean(accessToken && id) && all_admin_roles.includes(userRole ?? "") }, + ); +}; + +const useCustomerMutation = (action: (accessToken: string, input: TInput) => Promise) => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + return useMutation({ + mutationFn: (input) => { + if (!accessToken) throw new Error("Access token is required"); + return action(accessToken, input); + }, + onSuccess: () => + Promise.all([ + queryClient.invalidateQueries({ queryKey: ["get", "/customer/list"] }), + queryClient.invalidateQueries({ queryKey: ["get", "/customer/info"] }), + ]), + }); +}; + +export const useCreateCustomer = () => useCustomerMutation(createCustomer); +export const useUpdateCustomer = () => useCustomerMutation(updateCustomer); +export const useDeleteCustomers = () => useCustomerMutation(deleteCustomers); +export const useSetCustomerBlockedState = () => useCustomerMutation(setCustomerBlockedState); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index f4489625dfe..d80b88d8395 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -41,6 +41,7 @@ const LEGACY_PAGE_ROUTES: ReadonlyMap = new Map( agents: "agents", "router-settings": "router-settings", users: "users", + customers: "customers", teams: "teams", organizations: "organizations", }), diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 9a6ec334b56..649f9484088 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -34,6 +34,7 @@ import { Boxes, ChevronRight, Code2, + Contact, Database, ExternalLink, FileText, @@ -349,6 +350,7 @@ const menuGroups: MenuGroup[] = [ roles: [...projectReaderRoles], }, { key: "users", page: "users", label: "Internal Users", icon: , roles: all_admin_roles }, + { key: "customers", page: "customers", label: "Customers", icon: , roles: all_admin_roles }, { key: "organizations", page: "organizations", diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx index 47fb5643b22..4670293d7b8 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx @@ -47,6 +47,20 @@ describe("MCPServerSelector no-mcp-servers option", () => { setupMcpMocks(); }); + it("focuses the selector when its form label is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders( + <> + + + , + ); + + await user.click(screen.getByText("MCP Access")); + + expect(screen.getByLabelText("MCP Access")).toHaveFocus(); + }); + it("hides the No MCP Servers option by default", async () => { const user = userEvent.setup(); renderWithProviders( diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index 969db57462f..b98c7247d8d 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -6,6 +6,7 @@ import React from "react"; import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; interface MCPServerSelectorProps { + id?: string; onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; value?: { servers: string[]; @@ -24,6 +25,7 @@ interface MCPServerSelectorProps { const TOOLSET_PREFIX = "toolset:"; const MCPServerSelector: React.FC = ({ + id, onChange, value, className, @@ -107,6 +109,7 @@ const MCPServerSelector: React.FC = ({ return (
{ vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); }); + it("preserves an unrestricted stored grant until a tool permission is edited", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const tools = [ + { name: "list_records", description: "List records" }, + { name: "delete_record", description: "Delete a record" }, + ]; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([{ server_id: mockServerId, server_name: mockServerName }]); + vi.mocked(networking.listMCPTools).mockResolvedValue({ tools, error: false }); + renderWithProviders( + , + ); + + expect(await screen.findByText("delete_record")).toBeInTheDocument(); + await user.click(screen.getByRole("radio", { name: "Flat List" })); + + expect(screen.getByRole("checkbox", { name: "delete_record" })).toBeChecked(); + expect(onChange).not.toHaveBeenCalled(); + + await user.click(screen.getByRole("checkbox", { name: "delete_record" })); + + expect(onChange).toHaveBeenCalledExactlyOnceWith({ [mockServerId]: ["list_records"] }); + }); + it("should update tool permissions when user selects a tool", async () => { /** * Tests that clicking a tool checkbox calls onChange with updated permissions. diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index 7edeaedff2e..2e4ed053931 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -27,6 +27,7 @@ interface MCPToolPermissionsProps { toolPermissions: Record; onChange: (toolPermissions: Record) => void; disabled?: boolean; + initializeDefaultPermissions?: boolean; } const NO_SELECTION: readonly string[] = []; @@ -57,6 +58,7 @@ const MCPToolPermissions: React.FC = ({ toolPermissions, onChange, disabled = false, + initializeDefaultPermissions = true, }) => { const { data: allServers = [], @@ -110,6 +112,7 @@ const MCPToolPermissions: React.FC = ({ } else { const fetchedTools: MCPTool[] = response.tools || []; setServerTools((prev) => ({ ...prev, [serverId]: fetchedTools })); + if (!initializeDefaultPermissions) return; // Default only unrestricted direct servers to non-delete tools. // Read latest permissions from the ref to avoid clobbering concurrent results. diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index ae46729222a..8b4e440a617 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -53,7 +53,7 @@ export function ObjectPermissionsView({

Search tools

{searchTools.length === 0 ? (

- No restriction — all configured search tools are allowed for this team. + No restriction. All configured search tools are allowed

) : (

{searchTools.join(", ")}

diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx index e34359687d9..b30879dd7c7 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx @@ -464,6 +464,7 @@ describe("MCPServerPermissions", () => { />, ); await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive")); + expect(screen.getByText("No MCP servers. Access to all MCP servers is blocked")).toBeInTheDocument(); }); it("lists servers inherited from access groups, counts them, and names the group on hover", async () => { diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx index b4546d0b4dd..a2548a9f140 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -144,9 +144,7 @@ export function MCPServerPermissions({ {blocksAllMcpServers ? (
-

- No MCP servers — this key is blocked from all MCP servers, including its team's servers -

+

No MCP servers. Access to all MCP servers is blocked

) : grantsAllProxyMcpServers ? (
From bdefba9712b9a0db06d0f4c2f448141671046f8b Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 18:01:14 -0400 Subject: [PATCH 04/56] fix(ui): prevent stale customer MCP grants --- .../CustomersPanel.integration.test.tsx | 40 +++++ .../_components/customerPayload.test.ts | 24 +++ .../customers/_components/customerPayload.ts | 15 +- .../_components/useSaveCustomer.test.tsx | 160 ++++++++++++++++++ .../customers/_components/useSaveCustomer.ts | 10 +- .../view_users/user_info_view.test.tsx | 30 ++++ .../MCPToolPermissions.test.tsx | 30 ++++ .../MCPToolPermissions.tsx | 11 +- .../mcp_server_management/mcpEntitlement.ts | 16 +- .../src/components/page_metadata.ts | 1 + 10 files changed, 313 insertions(+), 24 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx index 31df2c883f6..3c3acc4d9fd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx @@ -318,6 +318,46 @@ describe("Customers management", () => { ); }); + it("shows a catalog failure when removing an MCP grant and keeps the customer unsaved", async () => { + const user = userEvent.setup(); + const directCustomer: EndUser = { + ...activeCustomer, + object_permission: { + object_permission_id: "permission-alice", + mcp_servers: ["server-1"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-1": ["read_document"] }, + agent_access_groups: [], + agents: [], + blocked_tools: [], + models: [], + search_tools: [], + skills: [], + vector_stores: [], + }, + }; + useQuery.mockReturnValue({ data: directCustomer, isLoading: false, isError: false }); + fetchServers.mockRejectedValue(new Error("Catalog unavailable")); + renderWithProviders( + <> + + + , + { searchParams: "customer=alice%40example.test&edit=true" }, + ); + + await user.click(await screen.findByRole("button", { name: "Clear all" })); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect( + await screen.findByText("MCP catalogs must finish loading before permission changes can be saved"), + ).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Save Changes" })).toBeInTheDocument(); + expect(screen.queryByText("Customer updated")).not.toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + }); + it("rejects an edit budget of zero because customer updates cannot apply it", async () => { const user = userEvent.setup(); renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts index 33d68688e2f..d5f8b85c30a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts @@ -79,6 +79,30 @@ describe("customer payload", () => { expect(buildCustomerPayload(values, null, [], [])).toEqual({ user_id: customer.user_id }); }); + it("drops generated tool grants when a newly selected server is removed before create", () => { + const values = { + ...toCustomerFormValues(null), + user_id: customer.user_id, + mcp_tool_permissions: { [server.server_id]: ["read"] }, + }; + + expect(buildCustomerPayload(values, null, [server], [])).toEqual({ user_id: customer.user_id }); + }); + + it("drops a new server's generated tool grants after returning to the stored selection", () => { + const stored = { ...customer, object_permission: { ...permission, mcp_access_groups: [], mcp_toolsets: [] } }; + const anotherServer = { ...server, server_id: "another-server" }; + const values = { + ...toCustomerFormValues(stored), + mcp_tool_permissions: { ...permission.mcp_tool_permissions, [anotherServer.server_id]: ["list"] }, + }; + + expect(buildCustomerPayload(values, stored, [server, anotherServer], [])).toEqual({ + user_id: customer.user_id, + models: [], + }); + }); + it("keeps zero budgets and trims optional text fields", () => { const values = { ...toCustomerFormValues(null), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts index de0a49d6eab..68ba8fbd75c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts @@ -59,9 +59,6 @@ const grantIdentity = (grant: McpEntitlementUpdate): string => { return JSON.stringify(normalized); }; -const selectionIdentity = (grant: McpEntitlementUpdate): string => - grantIdentity({ ...grant, mcp_tool_permissions: {} }); - const extractEditedMcpGrant = ( values: CustomerFormValues, storedGrant: McpEntitlementUpdate, @@ -103,19 +100,9 @@ const changedMcpGrant = ( servers: MCPServer[], toolsets: MCPToolset[], ): McpEntitlementUpdate | null => { - const selection = values.mcp_servers_and_groups; - const formGrant = { - mcp_servers: selection.servers, - mcp_access_groups: selection.accessGroups, - mcp_toolsets: selection.toolsets, - mcp_tool_permissions: values.mcp_tool_permissions, - }; const storedGrant = storedMcpGrant(existing); const storedIdentity = grantIdentity(storedGrant); - const grant = - selectionIdentity(formGrant) === selectionIdentity(storedGrant) - ? formGrant - : extractEditedMcpGrant(values, storedGrant, servers, toolsets); + const grant = extractEditedMcpGrant(values, storedGrant, servers, toolsets); return grant && grantIdentity(grant) !== storedIdentity ? grant : null; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.test.tsx new file mode 100644 index 00000000000..5fc62578019 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.test.tsx @@ -0,0 +1,160 @@ +import { renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import { toCustomerFormValues } from "./customerPayload"; +import { useSaveCustomer } from "./useSaveCustomer"; + +const mocks = vi.hoisted(() => ({ servers: vi.fn(), toolsets: vi.fn(), create: vi.fn(), update: vi.fn() })); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: mocks.servers })); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ useMCPToolsets: mocks.toolsets })); +vi.mock("@/app/(dashboard)/hooks/customers/useCustomerMutations", () => ({ + useCreateCustomer: () => ({ mutateAsync: mocks.create, isPending: false }), + useUpdateCustomer: () => ({ mutateAsync: mocks.update, isPending: false }), +})); + +const permission = { + object_permission_id: "permission-id", + mcp_servers: ["server-id"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-id": ["read"] }, + agent_access_groups: [], + agents: [], + blocked_tools: [], + models: [], + search_tools: [], + skills: [], + vector_stores: [], +}; +const customer: EndUser = { + user_id: "customer-id", + blocked: false, + spend: 0, + object_permission: permission, +}; +const server: MCPServer = { + server_id: "server-id", + created_at: "2026-01-01T00:00:00Z", + created_by: "admin", + updated_at: "2026-01-01T00:00:00Z", + updated_by: "admin", +}; +const toolset: MCPToolset = { + toolset_id: "toolset-id", + toolset_name: "Toolset", + tools: [{ server_id: server.server_id, tool_name: "read" }], +}; +const clearedGrant = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; + +beforeEach(() => { + vi.clearAllMocks(); + mocks.servers.mockReturnValue({ data: undefined, isSuccess: false }); + mocks.toolsets.mockReturnValue({ data: undefined, isSuccess: false }); + mocks.create.mockResolvedValue(customer); + mocks.update.mockResolvedValue(customer); +}); + +describe("useSaveCustomer", () => { + it.each([ + { state: "pending", data: undefined }, + { state: "failed", data: [server] }, + ])("refuses a server revocation while the catalog is $state", async ({ data }) => { + mocks.servers.mockReturnValue({ data, isSuccess: false }); + const { result } = renderHook(() => useSaveCustomer(customer)); + const values = { + ...toCustomerFormValues(customer), + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow("MCP catalogs must finish loading"); + expect(mocks.update).not.toHaveBeenCalled(); + }); + + it("saves an ordinary edit without catalogs and leaves unknown stored grants untouched", async () => { + const stored: EndUser = { + ...customer, + object_permission: { + ...permission, + mcp_servers: ["unknown-server"], + mcp_tool_permissions: { "unknown-server": ["read"] }, + }, + }; + const { result } = renderHook(() => useSaveCustomer(stored)); + + await result.current.saveCustomer({ ...toCustomerFormValues(stored), alias: "New alias" }); + + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ user_id: stored.user_id, alias: "New alias", models: [] }); + }); + + it("revokes a direct server once its catalog loads, without requiring an unused toolset catalog", async () => { + mocks.servers.mockReturnValue({ data: [server], isSuccess: true }); + const { result } = renderHook(() => useSaveCustomer(customer)); + const values = { + ...toCustomerFormValues(customer), + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + }; + + await result.current.saveCustomer(values); + + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ + user_id: customer.user_id, + models: [], + object_permission: clearedGrant, + }); + }); + + it("refuses a revocation with an empty final permission map until the server catalog loads", async () => { + const stored: EndUser = { ...customer, object_permission: { ...permission, mcp_tool_permissions: {} } }; + const { result } = renderHook(() => useSaveCustomer(stored)); + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow("MCP catalogs must finish loading"); + expect(mocks.update).not.toHaveBeenCalled(); + }); + + it("refuses a new customer's server grant until the server catalog loads", async () => { + const { result } = renderHook(() => useSaveCustomer(null)); + const values = { + ...toCustomerFormValues(null), + user_id: customer.user_id, + mcp_servers_and_groups: { servers: [server.server_id], accessGroups: [], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow("MCP catalogs must finish loading"); + expect(mocks.create).not.toHaveBeenCalled(); + }); + + it.each([ + { state: "pending", data: undefined }, + { state: "failed", data: [toolset] }, + ])("refuses a toolset revocation while its catalog is $state, then saves after it loads", async ({ data }) => { + const stored: EndUser = { + ...customer, + object_permission: { ...permission, mcp_servers: [], mcp_toolsets: [toolset.toolset_id] }, + }; + mocks.servers.mockReturnValue({ data: [server], isSuccess: true }); + mocks.toolsets.mockReturnValue({ data, isSuccess: false }); + const { result, rerender } = renderHook(() => useSaveCustomer(stored)); + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow("MCP catalogs must finish loading"); + expect(mocks.update).not.toHaveBeenCalled(); + mocks.toolsets.mockReturnValue({ data: [toolset], isSuccess: true }); + rerender(); + + await result.current.saveCustomer(values); + + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ + user_id: stored.user_id, + models: [], + object_permission: clearedGrant, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts index a0046ad8d25..a5e48a0423e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts @@ -5,12 +5,18 @@ import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolset import { buildCustomerPayload, type CustomerFormValues } from "./customerPayload"; export const useSaveCustomer = (existing: EndUser | null) => { - const { data: servers = [] } = useMCPServers(); - const { data: toolsets = [] } = useMCPToolsets(); + const { data: servers = [], isSuccess: serversReady } = useMCPServers(); + const { data: toolsets = [], isSuccess: toolsetsReady } = useMCPToolsets(); const create = useCreateCustomer(); const update = useUpdateCustomer(); const saveCustomer = (values: CustomerFormValues) => { const payload = buildCustomerPayload(values, existing, servers, toolsets); + const needsToolsets = + values.mcp_servers_and_groups.toolsets.length > 0 || (existing?.object_permission?.mcp_toolsets?.length ?? 0) > 0; + const catalogsReady = serversReady && (!needsToolsets || toolsetsReady); + if (payload.object_permission && !catalogsReady) { + return Promise.reject(new Error("MCP catalogs must finish loading before permission changes can be saved")); + } return existing ? update.mutateAsync(payload) : create.mutateAsync(payload); }; return { saveCustomer, isSaving: create.isPending || update.isPending }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx index 3eae1dbdfc5..06dec9e9614 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx @@ -3,6 +3,7 @@ import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; import UserInfoView from "./user_info_view"; import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; +import type { MCPServer } from "@/components/mcp_tools/types"; const mockTeamMemberAddCall = vi.fn(); const mockTeamMemberDeleteCall = vi.fn(); @@ -441,6 +442,12 @@ describe("UserInfoView", () => { }); describe("extractMcpEntitlement", () => { + const auditFields = { + created_at: "2026-01-01", + created_by: "test-admin", + updated_at: "2026-01-01", + updated_by: "test-admin", + }; const CATALOG = [ { server_id: "srv-1", server_name: "deploy_tracker", alias: "deploy" }, { server_id: "srv-2", server_name: "issue_tracker", alias: null }, @@ -494,6 +501,29 @@ describe("extractMcpEntitlement", () => { expect(result?.mcp_tool_permissions).toEqual({}); }); + it.each(["server_name", "alias"])("revokes an ID-keyed grant despite another server's matching %s", (field) => { + const catalog: MCPServer[] = [ + { ...auditFields, server_id: "srv-1", server_name: "First Server" }, + { ...auditFields, server_id: "srv-2", server_name: "Second Server", [field]: "srv-1" }, + ]; + const result = extractMcpEntitlement(form({ servers: ["srv-2"] }, { "srv-1": ["read"] }), catalog); + + expect(result?.mcp_tool_permissions).toEqual({}); + }); + + it.each(["server_name", "alias"])("resolves a selected ID before another server's matching %s", (field) => { + const catalog: MCPServer[] = [ + { ...auditFields, server_id: "srv-1", server_name: "First Server" }, + { ...auditFields, server_id: "srv-2", server_name: "Second Server", [field]: "srv-1" }, + ]; + const result = extractMcpEntitlement( + form({ servers: ["srv-1"] }, { "srv-1": ["read"], "srv-2": ["write"] }), + catalog, + ); + + expect(result?.mcp_tool_permissions).toEqual({ "srv-1": ["read"] }); + }); + it("prunes nothing when the server catalog has not loaded", () => { // Every key is unresolvable without the catalog, and pruning while under-informed is the // direction that widens. diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx index fdce2444c69..895385e5d02 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.test.tsx @@ -52,6 +52,36 @@ describe("MCPToolPermissions", () => { expect(onChange).toHaveBeenCalledExactlyOnceWith({ [mockServerId]: ["list_records"] }); }); + it("defaults a newly added edit server to non-delete tools without narrowing stored servers", async () => { + const onChange = vi.fn(); + const newServerId = "new-server"; + const tools = [ + { name: "list_records", description: "List records" }, + { name: "delete_record", description: "Delete a record" }, + ]; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId, server_name: mockServerName }, + { server_id: newServerId, server_name: "New Server" }, + ]); + vi.mocked(networking.listMCPTools).mockResolvedValue({ tools, error: false }); + const props = { + accessToken: mockAccessToken, + toolPermissions: {}, + onChange, + initializeDefaultPermissions: false, + }; + const { rerender } = renderWithProviders(); + + expect(await screen.findByText("delete_record")).toBeInTheDocument(); + expect(onChange).not.toHaveBeenCalled(); + + rerender(); + + await waitFor(() => { + expect(onChange).toHaveBeenCalledExactlyOnceWith({ [newServerId]: ["list_records"] }); + }); + }); + it("should update tool permissions when user selects a tool", async () => { /** * Tests that clicking a tool checkbox calls onChange with updated permissions. diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx index 2e4ed053931..2515dc159fe 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPToolPermissions.tsx @@ -16,6 +16,7 @@ import { emptyMcpAccessGroups, mcpAllowedToolsFor, mcpGrantsAllTools, + mcpServersForIdentifier, resolveEffectiveMcpServers, } from "./effectiveMcpServers"; @@ -72,6 +73,7 @@ const MCPToolPermissions: React.FC = ({ const [loadingTools, setLoadingTools] = useState>({}); const [toolErrors, setToolErrors] = useState>({}); const [viewModes, setViewModes] = useState>({}); + const initialSelectedServers = useRef(selectedServers); // Keep a ref to the latest toolPermissions so async fetch callbacks always // read the current value and do not overwrite sibling servers' results when @@ -112,7 +114,14 @@ const MCPToolPermissions: React.FC = ({ } else { const fetchedTools: MCPTool[] = response.tools || []; setServerTools((prev) => ({ ...prev, [serverId]: fetchedTools })); - if (!initializeDefaultPermissions) return; + if ( + !initializeDefaultPermissions && + initialSelectedServers.current.some((identifier) => + mcpServersForIdentifier(allServers, identifier).some((server) => server.server_id === serverId), + ) + ) { + return; + } // Default only unrestricted direct servers to non-delete tools. // Read latest permissions from the ref to avoid clobbering concurrent results. diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts index 66f59f5b71d..a6101b28a71 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts +++ b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts @@ -1,5 +1,6 @@ import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import { mcpServersForIdentifier } from "./effectiveMcpServers"; export interface McpEntitlementUpdate { mcp_servers: string[]; @@ -18,9 +19,6 @@ const asToolPermissions = (value: unknown): Record => { ); }; -const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): boolean => - server.server_id === identifier || server.server_name === identifier || server.alias === identifier; - /** * The `object_permission` a save sends, derived from what the editor currently shows. * @@ -35,8 +33,7 @@ const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): bool * which is the same membership the gateway resolves against. A selected toolset missing from the * catalog is unresolvable, so nothing is pruned in that save. * - * A tool-permission key may be a server id, a name or an alias: the gateway normalizes all three - * before looking up the allowlist, so an entry written by the API or by config can use any of them. + * A key matching a server id names only that server. Other keys fall back to names and aliases. * `allServers` is what resolves a key to its servers, plural: names and aliases are not unique, and * the gateway unions such a key into EVERY server answering to it, so the entry is kept while any * one of them is still granted. Resolving to the first match instead would make the outcome depend @@ -66,14 +63,19 @@ export const extractMcpEntitlement = ( .filter((toolset) => mcpToolsets.includes(toolset.toolset_id)) .flatMap((toolset) => toolset.tools.map((tool) => tool.server_id)), ); + const directServerIds = new Set( + mcpServers.flatMap((identifier) => + mcpServersForIdentifier(allServers, identifier).map((server) => server.server_id), + ), + ); const grants = (server: MCPServer): boolean => - mcpServers.some((identifier) => mcpServerMatchesIdentifier(server, identifier)) || + directServerIds.has(server.server_id) || (server.mcp_access_groups ?? []).some((group) => mcpAccessGroups.includes(group)) || toolsetServerIds.has(server.server_id); const grantsServerNamedBy = (permissionKey: string): boolean => { - const named = allServers.filter((candidate) => mcpServerMatchesIdentifier(candidate, permissionKey)); + const named = mcpServersForIdentifier(allServers, permissionKey); if (named.length === 0) return true; return named.some(grants); }; diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 72bed431421..17e8805e087 100644 --- a/ui/litellm-dashboard/src/components/page_metadata.ts +++ b/ui/litellm-dashboard/src/components/page_metadata.ts @@ -26,6 +26,7 @@ export const pageDescriptions: Record = { lens: "Review agent activity and investigate patterns with supporting evidence", "guardrails-monitor": "Monitor guardrail performance and view logs", users: "Manage internal user accounts and permissions", + customers: "Manage end-user spend, budgets, rate limits and model access", teams: "Create and manage teams for access control", organizations: "Manage organizations and their members", projects: "Manage projects within teams", From 6825f92806515f41319709496aca7bf569a6bdba Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 18:22:49 -0400 Subject: [PATCH 05/56] fix(proxy): enforce customer toolset server ceilings --- .../mcp_server/auth/user_api_key_auth_mcp.py | 21 +- .../auth/test_user_api_key_auth_mcp.py | 198 +++++++++++++++++- 2 files changed, 205 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 18b7c279606..934fa3fe565 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -3189,9 +3189,17 @@ class MCPRequestHandler: verbose_logger.debug("prisma_client is None") return [] - object_permission = await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + object_permission: Final = await MCPRequestHandler._get_end_user_object_permission( + user_api_key_auth, prisma_client + ) if object_permission is None: return [] + if SpecialMCPServerNames.no_mcp_servers.value in (object_permission.mcp_servers or []): + return [SpecialMCPServerNames.no_mcp_servers.value] + + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + object_permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) try: # Permission entries may be server_ids OR names/aliases — expand to ids. @@ -3199,7 +3207,9 @@ class MCPRequestHandler: global_mcp_server_manager, ) - direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permission.mcp_servers or []) + direct_mcp_servers: Final = global_mcp_server_manager.expand_permission_list( + object_permission.mcp_servers or [] + ) # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups( @@ -3212,9 +3222,10 @@ class MCPRequestHandler: global_mcp_server_manager.expand_tool_permissions(object_permission.mcp_tool_permissions).keys() ) - # Combine all lists - all_servers: Final = direct_mcp_servers + access_group_servers + tool_perm_servers - return list(MCPRequestHandler._server_grants_with_group_scope(set(all_servers), object_permission)) + all_servers: Final = frozenset( + (*direct_mcp_servers, *access_group_servers, *tool_perm_servers, *toolset_grants) + ) + return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, object_permission)) except Exception as e: verbose_logger.warning("Failed to get allowed MCP servers for end_user: %s", e) return [] diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 078fff0266b..9010b8584a6 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -10,12 +10,14 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers +from litellm.caching.dual_cache import DualCache from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, _agent_capped_servers, is_mcp_admitted_user_subject, ) +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, ProxyException, @@ -24,6 +26,8 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.agents import AgentCaller +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_toolset import MCPToolset @pytest.mark.parametrize("headers,expected", [ @@ -10535,26 +10539,32 @@ class TestEntitlementFaultSemantics: assert set(allowed) == {"srv1"} -async def _cache_with_end_user(end_user_id, *, mcp_tool_permissions=None, object_permission_id=None): +async def _cache_with_end_user( + end_user_id: str, + *, + mcp_tool_permissions: dict[str, list[str]] | None = None, + object_permission_id: str | None = None, + object_permission_override: LiteLLM_ObjectPermissionTable | None = None, +) -> DualCache: """A real DualCache already holding the end user row, so ``get_end_user_object`` answers from cache and no ``litellm.`` internal has to be patched. ``object_permission_id`` without a permission body models a row that NAMES an entitlement the DB then fails to serve.""" - from litellm.caching.dual_cache import DualCache from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key - cache = DualCache() + permission: Final = object_permission_override or ( + LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_tool_permissions=mcp_tool_permissions) + if mcp_tool_permissions + else None + ) + cache: Final = DualCache() await cache.async_set_cache( key=end_user_cache_key(end_user_id), value=LiteLLM_EndUserTable( user_id=end_user_id, blocked=False, - object_permission_id=object_permission_id or ("op-eu" if mcp_tool_permissions else None), - object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="op-eu", mcp_tool_permissions=mcp_tool_permissions - ) - if mcp_tool_permissions - else None, + object_permission_id=object_permission_id or (permission.object_permission_id if permission else None), + object_permission=permission, ), ) return cache @@ -10565,6 +10575,176 @@ class TestEndUserToolCeiling: """The end user (customer) level narrows the TOOLS axis exactly as it narrows the servers axis, so `object_permission.mcp_tool_permissions` on `/customer/new` is enforced, not just stored.""" + @pytest.mark.parametrize("require_defined", (False, True)) + @pytest.mark.parametrize( + "key_servers,expected_servers", ((("server-a", "server-b"), ("server-a",)), (("server-b",), ())) + ) + async def test_end_user_toolsets_cap_shared_key_servers( + self, require_defined: bool, key_servers: tuple[str, ...], expected_servers: tuple[str, ...] + ) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"server-a": ["read"]}) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=list(key_servers) + ), + ) + with ( + _entitlement_fault_globals(user_api_key_cache=cache), + patch( + "litellm.proxy.proxy_server.general_settings", {"require_end_user_mcp_access_defined": require_defined} + ), + ): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + allowed_tools: Final = await MCPRequestHandler.get_allowed_tools_for_server("server-a", auth) + read_allowed: Final = await MCPRequestHandler.is_tool_allowed_for_server("read", "server-a", auth) + write_allowed: Final = await MCPRequestHandler.is_tool_allowed_for_server("write", "server-a", auth) + assert set(allowed_servers) == set(expected_servers) + assert allowed_tools == ["read"] + assert read_allowed is True + assert write_allowed is False + + async def test_end_user_toolset_servers_union_with_existing_grants(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", + mcp_servers=["direct"], + mcp_access_groups=["readers"], + mcp_tool_permissions={"inline": ["read"]}, + mcp_toolsets=["toolset-a"], + ) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"toolset": ["read"]}) + manager: Final = MCPServerManager() + manager.config_mcp_servers = { + "group": MCPServer(server_id="group", name="group", transport="http", access_groups=["readers"]), + } + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + mcp_servers=["direct", "group", "inline", "toolset", "other"], + ), + ) + with ( + _entitlement_fault_globals(user_api_key_cache=cache), + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", manager), + ): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert set(allowed_servers) == {"direct", "group", "inline", "toolset"} + + @pytest.mark.parametrize( + "toolset_rows", + ((), (MCPToolset(toolset_id="toolset-a", toolset_name="Empty", tools=[]),)), + ids=("missing-toolset", "empty-toolset"), + ) + async def test_end_user_unresolvable_toolsets_deny_shared_key_servers( + self, toolset_rows: tuple[MCPToolset, ...] + ) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + prisma: Final = MagicMock() + prisma.configure_mock(**{"db.litellm_mcptoolsettable.find_many": AsyncMock(return_value=list(toolset_rows))}) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["server-a"]), + ) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(auth) + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == [] + + async def test_end_user_toolset_ceiling_reads_fresh_policy(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"old-server": ["read"]}) + prisma: Final = MagicMock() + prisma.configure_mock( + **{ + "writer_db.litellm_mcptoolsettable.find_many": AsyncMock( + return_value=[ + MCPToolset( + toolset_id="toolset-a", + toolset_name="Updated", + tools=[ + {"server_id": "new-server", "tool_name": "read"}, + ], + ), + ] + ), + } + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", + mcp_servers=["old-server", "new-server"], + ), + ).model_copy(update={"requires_fresh_policy": True}) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user(auth) + assert allowed_servers == ["new-server"] + + async def test_end_user_toolset_writer_fault_does_not_drop_ceiling(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="op-eu", mcp_toolsets=["toolset-a"]) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"old-server": ["read"]}) + prisma: Final = MagicMock() + prisma.configure_mock( + **{ + "writer_db.litellm_mcptoolsettable.find_many": AsyncMock( + side_effect=RuntimeError("writer unavailable") + ), + } + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=["other-server"] + ), + ).model_copy(update={"requires_fresh_policy": True}) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == [] + + async def test_end_user_mcp_opt_out_overrides_toolset_servers(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", + mcp_servers=[SpecialMCPServerNames.no_mcp_servers.value], + mcp_toolsets=["toolset-a"], + ) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + await cache.async_set_cache(key="toolset_perms:toolset-a", value={"server-a": ["read"]}) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["server-a"]), + ) + with _entitlement_fault_globals(user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == [] + + async def test_end_user_empty_permission_keeps_inherited_servers(self) -> None: + cache: Final = await _cache_with_end_user( + "eu-1", + object_permission_override=LiteLLM_ObjectPermissionTable(object_permission_id="op-eu"), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["server-a"]), + ) + with _entitlement_fault_globals(user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert allowed_servers == ["server-a"] + async def test_end_user_tool_permissions_intersect_key_tools(self): auth = _key_auth_reaching("srv1", tools=["tool_a", "tool_b"], end_user_id="eu-1") cache = await _cache_with_end_user("eu-1", mcp_tool_permissions={"srv1": ["tool_a"]}) From f7bac1777837cfe9f56323a83d609022311559fa Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 18:53:26 -0400 Subject: [PATCH 06/56] fix(ui): align customers with admin page layouts --- .../customers/_components/CustomerForm.tsx | 21 +++-- .../_components/CustomerInfoView.tsx | 88 +++++++++---------- .../CustomersPanel.integration.test.tsx | 36 ++++++++ .../customers/_components/CustomersPanel.tsx | 35 ++------ .../customers/_components/CustomersTable.tsx | 47 +++++++--- .../_components/CustomersTableColumns.tsx | 2 +- 6 files changed, 139 insertions(+), 90 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx index f10b2f06e45..562c05cdab6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx @@ -15,6 +15,7 @@ import { FieldGroup } from "@/components/ui/field"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useZodForm } from "@/lib/forms/useZodForm"; +import { cn } from "@/lib/cva.config"; import { toCustomerFormValues, type CustomerFormValues } from "./customerPayload"; const mcpSelectionShape = { @@ -247,13 +248,19 @@ export function CustomerForm({ onChange={(permissions) => form.setValue("mcp_tool_permissions", permissions)} /> -
- - +
+
+ + +
); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx index a6da69ef4fc..141b89ea9cd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx @@ -10,7 +10,7 @@ import CopyButton from "@/components/shared/CopyButton"; import { PageHeader, PageHeaderTitle } from "@/components/shared/PageHeader"; import { ModelsCell, MoneyCell } from "@/components/shared/table_cells"; import { Button } from "@/components/ui/button"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Card, CardContent } from "@/components/ui/card"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { DropdownMenu, @@ -72,39 +72,35 @@ function CustomerOverview({ customer, accessToken }: Pick
- - Budget Status - - -
- + +

Budget Status

+
+
+ +
+
+ of{" "} + + + +
+

Reset: {getBudgetDurationLabel(budget?.budget_duration)}

-
- of{" "} - - - -
-

Reset: {getBudgetDurationLabel(budget?.budget_duration)}

 - - Rate Limits - -
+

Rate Limits

+
{budget?.tpm_limit ?? "Unlimited"} {budget?.rpm_limit ?? "Unlimited"}
- - Models - -
+

Models

+
@@ -124,7 +120,7 @@ function CustomerSettings({ customer, accessToken }: Pick{customer.budget_id ?? "No budget"} - + {getBudgetDurationLabel(budget?.budget_duration)} @@ -205,7 +201,7 @@ export default function CustomerInfoView({
- + {customer.alias || customer.user_id} -
-
-

Customer Settings

- {canModify && !showEdit && ( - + + +
+

Customer Settings

+ {canModify && !showEdit && ( + + )} +
+ {showEdit ? ( + + ) : ( + )} -
- {showEdit ? ( - - ) : ( - - )} -
+ + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx index 3c3acc4d9fd..bc29f655418 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx @@ -110,6 +110,42 @@ describe("Customers management", () => { ).not.toBeInTheDocument(); }); + it.each([ + { maxBudget: 0, budgetLabel: "$0.00" }, + { maxBudget: null, budgetLabel: "Unlimited" }, + ])( + "distinguishes a $budgetLabel budget from zero spend in the list and details", + async ({ maxBudget, budgetLabel }) => { + const user = userEvent.setup(); + const customer: EndUser = { + ...activeCustomer, + spend: 0, + litellm_budget_table: { + max_budget: maxBudget, + tpm_limit: 600, + rpm_limit: 6, + created_at: "2026-01-01T00:00:00Z", + }, + }; + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [customer] : customer, + isLoading: false, + isError: false, + })); + renderWithProviders(); + + const row = screen.getByRole("row", { name: /alice@example.test/ }); + expect(within(row).getByText(budgetLabel)).toBeInTheDocument(); + await user.click(within(row).getByRole("button", { name: /Alice.*alice@example.test/ })); + expect(await screen.findByText("Budget Status")).toBeInTheDocument(); + expect(screen.getAllByText("$0.00")).toHaveLength(maxBudget === 0 ? 2 : 1); + expect(screen.getAllByText(budgetLabel)).toHaveLength(maxBudget === 0 ? 2 : 1); + await user.click(screen.getByRole("tab", { name: "Settings" })); + expect(await screen.findByText("Max Budget")).toBeInTheDocument(); + expect(screen.getByText(budgetLabel)).toBeInTheDocument(); + }, + ); + it("keeps MCP access optional while showing each customer status", async () => { const user = userEvent.setup(); renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx index 43306e508e2..4efb4b764ea 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx @@ -1,6 +1,6 @@ "use client"; -import { Contact, Plus, RefreshCw } from "lucide-react"; +import { Contact, Plus } from "lucide-react"; import { parseAsBoolean, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; import { useCallback, useState } from "react"; @@ -29,17 +29,7 @@ const CUSTOMER_VIEW_PARSERS = { customer_tab: parseAsStringLiteral(["overview", "settings"]).withDefault("overview"), }; -function CustomersHeader({ - canModify, - onCreate, - onRefresh, - isRefreshing, -}: { - canModify: boolean; - onCreate: () => void; - onRefresh: () => void; - isRefreshing: boolean; -}) { +function CustomersHeader({ canModify, onCreate }: { canModify: boolean; onCreate: () => void }) { return ( @@ -47,18 +37,14 @@ function CustomersHeader({ Customers Manage end-user spend, budgets, rate limits, model and MCP access - - - {canModify && ( + {canModify && ( + - )} - + + )} ); } @@ -142,16 +128,13 @@ export default function CustomersPanel() { /> ) : ( <> - setIsCreateOpen(true)} - onRefresh={() => void customers.refetch()} - isRefreshing={customers.isFetching} - /> + setIsCreateOpen(true)} /> void customers.refetch()} + isRefreshing={customers.isFetching} canModify={canModify} premiumUser={premiumUser} onView={openCustomer} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx index f0f5e71368b..fdc0b2aae8c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx @@ -1,11 +1,12 @@ "use client"; -import { Contact } from "lucide-react"; +import { Contact, RefreshCw, Search } from "lucide-react"; import { useMemo } from "react"; import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; import { DataTable, DataTableViewOptions, useUrlTableState } from "@/components/shared/DataTable"; import { Input } from "@/components/ui/input"; +import { Button } from "@/components/ui/button"; import { getCustomerTableColumns, type CustomerTableActions } from "./CustomersTableColumns"; @@ -22,6 +23,8 @@ type CustomersTableProps = CustomerTableActions & { data: EndUser[]; isLoading: boolean; isError: boolean; + onRefresh: () => void; + isRefreshing: boolean; }; function EmptyCustomers() { @@ -34,7 +37,14 @@ function EmptyCustomers() { ); } -export default function CustomersTable({ data, isLoading, isError, ...actions }: CustomersTableProps) { +export default function CustomersTable({ + data, + isLoading, + isError, + onRefresh, + isRefreshing, + ...actions +}: CustomersTableProps) { const { search, setSearch, sorting, onSortingChange, pagination, onPaginationChange } = useUrlTableState(TABLE_STATE_OPTIONS); const { canModify, premiumUser, onView, onEdit, onDelete, onSetBlocked } = actions; @@ -57,15 +67,30 @@ export default function CustomersTable({ data, isLoading, isError, ...actions }: getRowId={(customer) => customer.user_id} defaultColumnVisibility={DEFAULT_COLUMN_VISIBILITY} toolbar={(table) => ( -
- setSearch(event.target.value)} - /> - +
+
+ + setSearch(event.target.value)} + /> +
+
+ + +
)} paginationMode="client" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx index 775273ad43a..7ce1146aa9b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx @@ -180,7 +180,7 @@ export function getCustomerTableColumns(actions: CustomerTableActions): ColumnDe header: ({ column }) => , size: 140, cell: ({ row }) => ( - + ), }, { From b1e0be1f997599d4c2ac32108dc21dbcc78266ff Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 18:53:55 -0400 Subject: [PATCH 07/56] fix(proxy): deny MCP access when a customer's server ceiling cannot be read --- .../mcp_server/auth/user_api_key_auth_mcp.py | 5 +-- .../auth/test_user_api_key_auth_mcp.py | 31 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 934fa3fe565..f36a12b58ef 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -3227,8 +3227,9 @@ class MCPRequestHandler: ) return list(MCPRequestHandler._server_grants_with_group_scope(all_servers, object_permission)) except Exception as e: - verbose_logger.warning("Failed to get allowed MCP servers for end_user: %s", e) - return [] + raise UnloadableEntitlementError( + f"end user {user_api_key_auth.end_user_id!r} has an object permission whose MCP servers could not be read" + ) from e @staticmethod async def _get_user_object_permission( diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 9010b8584a6..88d5f0cd387 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -10714,6 +10714,37 @@ class TestEndUserToolCeiling: allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) assert allowed_servers == [] + async def test_end_user_access_group_writer_fault_does_not_drop_toolset_ceiling(self) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="op-eu", mcp_toolsets=["toolset-a"], mcp_access_groups=["group-x"] + ) + cache: Final = await _cache_with_end_user("eu-1", object_permission_override=permission) + prisma: Final = MagicMock() + prisma.configure_mock( + **{ + "writer_db.litellm_mcptoolsettable.find_many": AsyncMock( + return_value=[ + MCPToolset( + toolset_id="toolset-a", + toolset_name="Read only", + tools=[{"server_id": "server-a", "tool_name": "read"}], + ), + ] + ), + "writer_db.litellm_mcpservertable.find_many": AsyncMock(side_effect=RuntimeError("writer unavailable")), + } + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="eu-1", + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-key", mcp_servers=["server-a", "server-b"] + ), + ).model_copy(update={"requires_fresh_policy": True}) + with _entitlement_fault_globals(prisma_client=prisma, user_api_key_cache=cache): + allowed_servers: Final = await MCPRequestHandler.get_allowed_mcp_servers(auth) + assert "server-b" not in allowed_servers + async def test_end_user_mcp_opt_out_overrides_toolset_servers(self) -> None: permission: Final = LiteLLM_ObjectPermissionTable( object_permission_id="op-eu", From 13a9e88608206e5449e359042c90fbeee0d95e8e Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 19:37:05 -0400 Subject: [PATCH 08/56] ci(test-unit): run the MCP auth unit tests in the proxy-auth shard --- .github/workflows/test-unit.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 23a340cd436..2e2de270745 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -457,6 +457,7 @@ jobs: - shard: proxy-auth test-path: |- tests/unit/proxy/auth + tests/unit/proxy/_experimental/mcp_server/auth --ignore=tests/unit/proxy/auth/test_auth_checks.py --ignore=tests/unit/proxy/auth/test_user_api_key_auth.py --ignore=tests/unit/proxy/auth/test_default_end_user_budget_simple.py From df17e4949cdcb637752ab8493388833a1b1a1504 Mon Sep 17 00:00:00 2001 From: lewis Date: Sun, 11 Oct 2026 20:06:59 -0400 Subject: [PATCH 09/56] fix(ui): clarify customer budget selection --- .../customers/_components/CustomerForm.tsx | 24 ++++++++----------- .../CustomersPanel.integration.test.tsx | 4 ++-- .../src/components/page_metadata.ts | 2 +- 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx index 562c05cdab6..ab85dcac7ee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx @@ -80,9 +80,14 @@ export function CustomerForm({ defaultValues: initialValues, }); const hasLinkedBudget = isEdit && Boolean(initialValues.budget_id); - const budgetDescription = isEdit - ? "Leaving both budget fields empty keeps the current budget" - : "Choose a budget or enter a maximum below"; + const budgetFieldCopy = hasLinkedBudget + ? { emptyOption: "Keep current budget", description: "Select another budget to change the link" } + : { + emptyOption: "No budget", + description: isEdit + ? "Leaving both budget fields empty keeps the current budget" + : "Choose a budget or enter a maximum below", + }; const handleSubmit = form.handleSubmit(async (values) => { if (hasLinkedBudget && !values.budget_id && values.max_budget.trim()) { form.setError("max_budget", { @@ -94,7 +99,7 @@ export function CustomerForm({ }); const selectedBudgetId = form.watch("budget_id"); const budgetOptions = [ - { value: "", label: "No budget" }, + { value: "", label: budgetFieldCopy.emptyOption }, ...(selectedBudgetId && !budgets.some((budget) => budget.budget_id === selectedBudgetId) ? [{ value: selectedBudgetId, label: selectedBudgetId }] : []), @@ -126,16 +131,7 @@ export function CustomerForm({ {({ ref, ...control }) => } - + {({ id, value, onChange }) => ( setSearch(event.target.value)} - /> -
-
- + +
+ {canModify && selectedIds.length > 0 && ( +
+ {selectedIds.length} selected + void handleBulkSetBlocked(selectedIds, true)} + /> + void handleBulkSetBlocked(selectedIds, false)} + /> + +
+ )} + - - - -
-
- )} + {({ get, set }) => + CUSTOMER_FILTER_FIELDS.map((field) => ( + + + + )) + } + + + ); + }} + filterMode="client" + columnFilters={columnFilters} + onColumnFiltersChange={(updater) => { + setRowSelection({}); + onColumnFiltersChange(updater); + }} paginationMode="client" pagination={pagination} onPaginationChange={onPaginationChange} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.tsx new file mode 100644 index 00000000000..c3bc3e7a46a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.tsx @@ -0,0 +1,55 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; + +import { CustomerResourcesCell } from "./CustomersTableColumns"; + +const customer: EndUser = { + user_id: "customer-id", + spend: 0, + blocked: false, + models: ["model-a", "model-b"], + object_permission: { + object_permission_id: "permission-id", + agent_access_groups: [], + agents: [], + blocked_tools: [], + mcp_servers: ["server-a"], + mcp_access_groups: ["group-a", "group-b"], + mcp_toolsets: ["toolset-a", "toolset-b", "toolset-c"], + mcp_tool_permissions: { "server-b": ["read"] }, + models: [], + search_tools: [], + vector_stores: [], + }, +}; + +describe("CustomerResourcesCell", () => { + it("shows separately labeled counts for models and each MCP grant type", () => { + render(); + expect(screen.getByLabelText("2 models")).toHaveTextContent("2"); + expect(screen.getByLabelText("2 MCP servers")).toHaveTextContent("2"); + expect(screen.getByLabelText("2 MCP access groups")).toHaveTextContent("2"); + expect(screen.getByLabelText("3 MCP toolsets")).toHaveTextContent("3"); + }); + + it("shows All for unrestricted models and excludes unsupported MCP sentinel counts", () => { + render( + , + ); + expect(screen.getByLabelText("All models")).toHaveTextContent("All"); + expect(screen.getByLabelText("0 MCP servers")).toHaveTextContent("0"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx index 7ce1146aa9b..30fc04c0ddf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx @@ -1,7 +1,7 @@ "use client"; import type { ColumnDef } from "@tanstack/react-table"; -import { Ban, Copy, MoreHorizontal, Pencil, ShieldCheck, Trash2 } from "lucide-react"; +import { Ban, Copy, Layers, MoreHorizontal, Pencil, Server, ShieldCheck, Trash2, Users, Wrench } from "lucide-react"; import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; @@ -20,9 +20,12 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp import { cn } from "@/lib/cva.config"; import { copyToClipboard } from "@/utils/dataUtils"; +import { getCustomerResourceCounts } from "./customerResources"; + export type CustomerTableActions = { canModify: boolean; premiumUser: boolean; + isBlocking: boolean; onView: (customerId: string) => void; onEdit: (customerId: string) => void; onDelete: (customer: EndUser) => void; @@ -60,7 +63,7 @@ export function CustomerBlockAction({ }) { const action = ( onSetBlocked(customer)} > @@ -68,11 +71,15 @@ export function CustomerBlockAction({ {customer.blocked ? "Unblock" : "Block"} ); - if (customer.blocked || premiumUser) return action; + if (premiumUser) return action; return ( - }> + + } + > {action} @@ -100,6 +107,46 @@ export function describeMcpGrant(permission: EndUser["object_permission"]): stri return parts.join(", ") || "No restriction"; } +const RESOURCE_ITEMS = [ + { key: "models", label: "models", icon: Layers, className: "bg-info/10 text-info ring-sky-600/20" }, + { key: "servers", label: "MCP servers", icon: Server, className: "bg-success/10 text-success ring-emerald-600/20" }, + { + key: "accessGroups", + label: "MCP access groups", + icon: Users, + className: + "bg-violet-50 text-violet-700 ring-violet-600/20 dark:bg-violet-950 dark:text-violet-300 dark:ring-violet-400/30", + }, + { key: "toolsets", label: "MCP toolsets", icon: Wrench, className: "bg-warning/10 text-warning ring-amber-600/20" }, +] as const; + +export function CustomerResourcesCell({ customer }: { customer: EndUser }) { + const counts = getCustomerResourceCounts(customer); + return ( +
+ {RESOURCE_ITEMS.map(({ key, label, icon: Icon, className }) => { + const count = counts[key]; + const countLabel = count === 1 ? label.slice(0, -1) : label; + const title = count === "all" ? `All ${label}` : `${count} ${countLabel}`; + return ( + + + ); + })} +
+ ); +} + function CustomerRowActions({ customer, canModify, @@ -107,6 +154,7 @@ function CustomerRowActions({ onEdit, onDelete, onSetBlocked, + isBlocking, }: CustomerTableActions & { customer: EndUser }) { return ( @@ -123,7 +171,12 @@ function CustomerRowActions({ Edit - + )} void copyToClipboard(customer.user_id, "Customer ID copied")}> @@ -165,6 +218,14 @@ export function getCustomerTableColumns(actions: CustomerTableActions): ColumnDe /> ), }, + { + id: "resources", + meta: { title: "Resources", skeleton: "chips" }, + header: "Resources", + size: 250, + enableSorting: false, + cell: ({ row }) => , + }, { id: "spend", accessorKey: "spend", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.test.ts new file mode 100644 index 00000000000..924ead7901c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.test.ts @@ -0,0 +1,167 @@ +import Papa from "papaparse"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +import { buildCustomersCsv, buildCustomersCsvRows, exportCustomersToCsv } from "./customerCsvExport"; + +const customer = (overrides: Partial = {}): EndUser => ({ + user_id: "customer-1", + blocked: false, + spend: 0, + ...overrides, +}); + +afterEach(() => { + vi.unstubAllGlobals(); + vi.useRealTimers(); +}); + +describe("buildCustomersCsvRows", () => { + it("exports customer identity, status, budget, model access, and MCP grants", () => { + const customerFields: Partial = { + user_id: "customer-42", + alias: "finance", + blocked: true, + spend: 12.5, + budget_id: "budget-42", + litellm_budget_table: { + created_at: "2026-01-01T00:00:00Z", + max_budget: 250, + soft_budget: 200, + budget_duration: "30d", + budget_reset_at: "2026-02-01T00:00:00Z", + tpm_limit: 1000, + rpm_limit: 50, + tpd_limit: 20000, + max_parallel_requests: 10, + }, + models: ["model-1", "model-2"], + allowed_model_region: "eu", + default_model: "model-1", + object_permission: { + object_permission_id: "permission-42", + models: [], + mcp_servers: ["server-1", "server-2"], + mcp_access_groups: ["group-1"], + mcp_toolsets: ["toolset-1"], + agents: [], + agent_access_groups: [], + blocked_tools: [], + search_tools: [], + vector_stores: [], + }, + }; + const [row] = buildCustomersCsvRows([customer(customerFields)]); + + const expectedRow = { + "Customer Alias": "finance", + "Customer ID": "customer-42", + Status: "Blocked", + "Spend (USD)": 12.5, + "Budget ID": "budget-42", + "Max Budget (USD)": 250, + "Soft Budget (USD)": 200, + "Budget Duration": "30d", + "Budget Reset At": "2026-02-01T00:00:00Z", + "TPM Limit": 1000, + "RPM Limit": 50, + "TPD Limit": 20000, + "Max Parallel Requests": 10, + Models: "model-1, model-2", + "MCP Servers": "server-1, server-2", + "MCP Access Groups": "group-1", + "MCP Toolsets": "toolset-1", + "Allowed Model Region": "eu", + "Default Model": "model-1", + }; + expect(row).toEqual(expectedRow); + }); + + it("keeps zero limits and spend while rendering missing optional values as empty cells", () => { + const missingFields: Partial = { + alias: null, + budget_id: null, + litellm_budget_table: null, + object_permission: null, + }; + const [minimal, zeroBudget] = buildCustomersCsvRows([ + customer(missingFields), + customer({ + user_id: "zero-budget", + litellm_budget_table: { + created_at: "2026-01-01T00:00:00Z", + budget_id: "nested-budget", + max_budget: 0, + tpm_limit: 0, + rpm_limit: 0, + }, + }), + ]); + + expect(minimal.Status).toBe("Active"); + expect(minimal["Spend (USD)"]).toBe(0); + expect(Object.entries(minimal).filter(([key]) => !["Customer ID", "Status", "Spend (USD)"].includes(key))).toEqual( + Object.keys(minimal) + .filter((key) => !["Customer ID", "Status", "Spend (USD)"].includes(key)) + .map((key) => [key, ""]), + ); + expect(zeroBudget["Budget ID"]).toBe("nested-budget"); + expect(zeroBudget["Max Budget (USD)"]).toBe(0); + expect(zeroBudget["TPM Limit"]).toBe(0); + expect(zeroBudget["RPM Limit"]).toBe(0); + }); +}); + +describe("buildCustomersCsv", () => { + it("round trips quotes, commas, and line breaks without splitting customer rows", () => { + const csv = buildCustomersCsv([ + customer({ alias: 'sales, "east"\nregion', models: ["model-1", "model-2"] }), + customer({ user_id: "customer-2" }), + ]); + const parsed = Papa.parse>(csv, { header: true }); + + expect(parsed.errors).toEqual([]); + expect(parsed.data).toHaveLength(2); + expect(parsed.data[0]["Customer Alias"]).toBe('sales, "east"\nregion'); + expect(parsed.data[0].Models).toBe("model-1, model-2"); + expect(parsed.data[1]["Customer ID"]).toBe("customer-2"); + }); + + it.each(["=SUM(A1:A9)", "+cmd", "-cmd", "@cmd", "\tcmd", "\rcmd"])("neutralizes formula-leading text %j", (alias) => { + const csv = buildCustomersCsv([customer({ alias })]); + const parsed = Papa.parse>(csv, { header: true }); + + expect(parsed.data[0]["Customer Alias"]).toBe(`'${alias}`); + }); +}); + +describe("exportCustomersToCsv", () => { + it("downloads the supplied customers, returns their count, and cleans up the download URL", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-01-15T12:00:00Z")); + const a = { href: "", download: "", click: vi.fn() }; + const createObjectURL = vi.fn().mockReturnValue("blob:customers-export"); + const revokeObjectURL = vi.fn(); + const createElement = vi.fn().mockReturnValue(a); + const appendChild = vi.fn(); + const removeChild = vi.fn(); + vi.stubGlobal("window", { URL: { createObjectURL, revokeObjectURL } }); + vi.stubGlobal("document", { createElement, body: { appendChild, removeChild } }); + + expect(exportCustomersToCsv([customer({ user_id: "visible-1" }), customer({ user_id: "visible-2" })])).toBe(2); + + const [blob] = createObjectURL.mock.calls[0] as [Blob]; + expect(blob.type).toBe("text/csv;charset=utf-8;"); + expect( + Papa.parse>(await blob.text(), { header: true }).data.map((row) => row["Customer ID"]), + ).toEqual(["visible-1", "visible-2"]); + expect(createElement).toHaveBeenCalledWith("a"); + expect(a.href).toBe("blob:customers-export"); + expect(a.download).toBe("customers_export_2026-01-15.csv"); + expect(appendChild).toHaveBeenCalledWith(a); + expect(a.click).toHaveBeenCalledOnce(); + expect(removeChild).toHaveBeenCalledWith(a); + expect(revokeObjectURL).toHaveBeenCalledWith("blob:customers-export"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.ts new file mode 100644 index 00000000000..9181f99462a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.ts @@ -0,0 +1,48 @@ +import Papa from "papaparse"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +const cell = (value: string | number | null | undefined): string | number => value ?? ""; + +export const buildCustomersCsvRows = (customers: EndUser[]): Record[] => + customers.map((customer) => { + const budget = customer.litellm_budget_table; + const permission = customer.object_permission; + return { + "Customer Alias": cell(customer.alias), + "Customer ID": customer.user_id, + Status: customer.blocked ? "Blocked" : "Active", + "Spend (USD)": customer.spend, + "Budget ID": cell(customer.budget_id ?? budget?.budget_id), + "Max Budget (USD)": cell(budget?.max_budget), + "Soft Budget (USD)": cell(budget?.soft_budget), + "Budget Duration": cell(budget?.budget_duration), + "Budget Reset At": cell(budget?.budget_reset_at), + "TPM Limit": cell(budget?.tpm_limit), + "RPM Limit": cell(budget?.rpm_limit), + "TPD Limit": cell(budget?.tpd_limit), + "Max Parallel Requests": cell(budget?.max_parallel_requests), + Models: (customer.models ?? []).join(", "), + "MCP Servers": (permission?.mcp_servers ?? []).join(", "), + "MCP Access Groups": (permission?.mcp_access_groups ?? []).join(", "), + "MCP Toolsets": (permission?.mcp_toolsets ?? []).join(", "), + "Allowed Model Region": cell(customer.allowed_model_region), + "Default Model": cell(customer.default_model), + }; + }); + +export const buildCustomersCsv = (customers: EndUser[]): string => + Papa.unparse(buildCustomersCsvRows(customers), { escapeFormulae: true }); + +export const exportCustomersToCsv = (customers: EndUser[]): number => { + const blob = new Blob([buildCustomersCsv(customers)], { type: "text/csv;charset=utf-8;" }); + const url = window.URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = `customers_export_${new Date().toISOString().split("T")[0]}.csv`; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + window.URL.revokeObjectURL(url); + return customers.length; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.test.ts new file mode 100644 index 00000000000..590f3d2bdda --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.test.ts @@ -0,0 +1,85 @@ +import { describe, expect, it } from "vitest"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; + +import { customerMatchesFilter, customerMatchesSearch, formatCustomerFilterValue } from "./customerFilters"; + +const customer: EndUser = { user_id: "customer-1", alias: "Research", blocked: false, spend: 0 }; +const permission: NonNullable = { + object_permission_id: "permission-1", + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: {}, + agent_access_groups: [], + agents: [], + models: [], + blocked_tools: [], + search_tools: [], + vector_stores: [], +}; + +describe("customer filters", () => { + it("searches ID and alias case-insensitively and trims the search", () => { + expect(customerMatchesSearch(customer, " CUSTOMER-1 ")).toBe(true); + expect(customerMatchesSearch(customer, "research")).toBe(true); + expect(customerMatchesSearch({ ...customer, alias: null }, "research")).toBe(false); + expect(customerMatchesSearch(customer, " ")).toBe(true); + }); + + it.each([false, true])("filters both blocked and active for blocked=%s", (blocked) => { + const row = { ...customer, blocked }; + expect(customerMatchesFilter(row, "blocked", "blocked")).toBe(blocked); + expect(customerMatchesFilter(row, "blocked", "active")).toBe(!blocked); + }); + + it.each([0, 25, null, undefined])("uses a monetary limit for max_budget=%s", (max_budget) => { + const row = { ...customer, litellm_budget_table: { created_at: "2026-01-01T00:00:00Z", max_budget } }; + expect(customerMatchesFilter(row, "max_budget", "yes")).toBe(max_budget != null); + expect(customerMatchesFilter(row, "max_budget", "no")).toBe(max_budget == null); + }); + + it("treats absent budgets and an unlimited budget link as no monetary budget", () => { + expect(customerMatchesFilter(customer, "max_budget", "no")).toBe(true); + expect(customerMatchesFilter({ ...customer, budget_id: "budget-id" }, "max_budget", "yes")).toBe(false); + }); + + it.each([ + { mcp_servers: ["server-1"] }, + { mcp_access_groups: ["group-1"] }, + { mcp_toolsets: ["toolset-1"] }, + { mcp_tool_permissions: { "server-1": ["read"] } }, + ])("detects configured MCP grants %j", (grant) => { + const row = { ...customer, object_permission: { ...permission, ...grant } }; + expect(customerMatchesFilter(row, "mcp_access", "yes")).toBe(true); + expect(customerMatchesFilter(row, "mcp_access", "no")).toBe(false); + }); + + it("does not mistake absent, empty, or explicitly denied MCP grants for access", () => { + const denied = { ...permission, mcp_servers: [NO_MCP_SERVERS_SENTINEL], mcp_access_groups: ["group-1"] }; + for (const object_permission of [ + undefined, + null, + permission, + denied, + { ...permission, mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL] }, + ]) { + expect(customerMatchesFilter({ ...customer, object_permission }, "mcp_access", "yes")).toBe(false); + expect(customerMatchesFilter({ ...customer, object_permission }, "mcp_access", "no")).toBe(true); + } + }); + + it("ignores invalid URL filter values and unrelated columns", () => { + expect(customerMatchesFilter(customer, "blocked", "unexpected")).toBe(true); + expect(customerMatchesFilter(customer, "max_budget", false)).toBe(true); + expect(customerMatchesFilter(customer, "mcp_access", "unexpected")).toBe(true); + expect(customerMatchesFilter(customer, "unrelated", "yes")).toBe(true); + }); + + it("formats filter chips with the labels shown in the drawer", () => { + expect(formatCustomerFilterValue("blocked", "active")).toBe("Active"); + expect(formatCustomerFilterValue("max_budget", "yes")).toBe("Has budget"); + expect(formatCustomerFilterValue("mcp_access", "no")).toBe("No MCP access"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.ts new file mode 100644 index 00000000000..bce9adf696d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.ts @@ -0,0 +1,54 @@ +import type { FilterFn } from "@tanstack/react-table"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +import { customerHasMcpAccess } from "./customerResources"; + +export const CUSTOMER_FILTER_FIELDS = [ + { + id: "blocked", + label: "Status", + options: [ + { value: "active", label: "Active" }, + { value: "blocked", label: "Blocked" }, + ], + }, + { + id: "max_budget", + label: "Budget", + options: [ + { value: "yes", label: "Has budget" }, + { value: "no", label: "No budget" }, + ], + }, + { + id: "mcp_access", + label: "MCP access", + options: [ + { value: "yes", label: "Has MCP access" }, + { value: "no", label: "No MCP access" }, + ], + }, +] as const; + +export const customerMatchesSearch = (customer: EndUser, search: string): boolean => + `${customer.user_id} ${customer.alias ?? ""}`.toLowerCase().includes(search.trim().toLowerCase()); + +export const customerMatchesFilter = (customer: EndUser, columnId: string, value: unknown): boolean => { + if (columnId === "blocked") { + if (value === "blocked") return Boolean(customer.blocked); + if (value === "active") return !customer.blocked; + return true; + } + if (value !== "yes" && value !== "no") return true; + if (columnId === "max_budget") return (customer.litellm_budget_table?.max_budget != null) === (value === "yes"); + if (columnId === "mcp_access") return customerHasMcpAccess(customer) === (value === "yes"); + return true; +}; + +export const customerFilterFn: FilterFn = (row, columnId, value: unknown) => + customerMatchesFilter(row.original, columnId, value); + +export const formatCustomerFilterValue = (columnId: string, value: unknown): string => + CUSTOMER_FILTER_FIELDS.find((field) => field.id === columnId)?.options.find((option) => option.value === value) + ?.label ?? String(value); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.test.ts new file mode 100644 index 00000000000..48c5b7a5b2b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.test.ts @@ -0,0 +1,137 @@ +import { describe, expect, it } from "vitest"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; + +import { customerHasMcpAccess, getCustomerResourceCounts, type CustomerResourceCounts } from "./customerResources"; + +const emptyPermission: NonNullable = { + object_permission_id: "permission-id", + agent_access_groups: [], + agents: [], + blocked_tools: [], + mcp_servers: [], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: {}, + models: [], + search_tools: [], + vector_stores: [], +}; + +const customer: EndUser = { user_id: "customer-id", spend: 0, blocked: false }; +const noMcpCounts: CustomerResourceCounts = { models: "all", servers: 0, accessGroups: 0, toolsets: 0 }; + +describe("getCustomerResourceCounts", () => { + it("counts distinct model and MCP grants", () => { + const configured: EndUser = { + ...customer, + models: ["model-a", "model-b", "model-a"], + object_permission: { + ...emptyPermission, + mcp_servers: ["server-a", "server-b", "server-a"], + mcp_access_groups: ["group-a", "group-b", "group-a"], + mcp_toolsets: ["toolset-a", "toolset-a"], + }, + }; + const expectedCounts = { models: 2, servers: 2, accessGroups: 2, toolsets: 1 }; + expect(getCustomerResourceCounts(configured)).toEqual(expectedCounts); + }); + + it.each([undefined, [], ["all-proxy-models"], ["model-a", "all-proxy-models"], ["*"]])( + "labels unrestricted models as all (%j)", + (models) => { + expect(getCustomerResourceCounts({ ...customer, models }).models).toBe("all"); + }, + ); + + it("counts tool-permission servers without counting a server twice", () => { + const configured = { + ...customer, + object_permission: { + ...emptyPermission, + mcp_servers: ["server-a"], + mcp_tool_permissions: { "server-a": ["read", "write"], "server-b": ["read"] }, + }, + }; + expect(getCustomerResourceCounts(configured).servers).toBe(2); + expect(customerHasMcpAccess(configured)).toBe(true); + }); + + it("excludes the unsupported all-MCP sentinel from customer server counts", () => { + const configured = { + ...customer, + object_permission: { ...emptyPermission, mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL, "server-a"] }, + }; + expect(getCustomerResourceCounts(configured).servers).toBe(1); + expect(customerHasMcpAccess(configured)).toBe(true); + }); + + it("does not claim MCP access from unsupported all-MCP sentinel-only grants", () => { + const configured = { + ...customer, + object_permission: { + ...emptyPermission, + mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL], + mcp_tool_permissions: { [ALL_PROXY_MCP_SERVERS_SENTINEL]: ["read"] }, + }, + }; + expect(getCustomerResourceCounts(configured).servers).toBe(0); + expect(customerHasMcpAccess(configured)).toBe(false); + }); + + it("zeroes every MCP grant when no-MCP overrides other grants", () => { + const configured = { + ...customer, + object_permission: { + ...emptyPermission, + mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL, "server-a"], + mcp_access_groups: ["group-a"], + mcp_toolsets: ["toolset-a"], + mcp_tool_permissions: { "server-a": ["read"] }, + }, + }; + expect(getCustomerResourceCounts(configured)).toEqual(noMcpCounts); + expect(customerHasMcpAccess(configured)).toBe(false); + }); + + it("excludes permission-table models and resources not enforced for customers", () => { + const configured = { + ...customer, + models: ["customer-model"], + object_permission: { + ...emptyPermission, + models: ["permission-model-a", "permission-model-b"], + vector_stores: ["store-id"], + agents: ["agent-id"], + agent_access_groups: ["agent-group-id"], + search_tools: ["search-tool-id"], + skills: ["skill-id"], + }, + }; + expect(getCustomerResourceCounts(configured)).toEqual({ ...noMcpCounts, models: 1 }); + expect(customerHasMcpAccess(configured)).toBe(false); + }); +}); + +describe("customerHasMcpAccess", () => { + it.each([ + undefined, + null, + emptyPermission, + { ...emptyPermission, mcp_servers: null, mcp_access_groups: null, mcp_toolsets: null }, + ])("treats an empty customer-level grant as not configured (%j)", (object_permission) => { + const unconfigured = { ...customer, object_permission }; + expect(customerHasMcpAccess(unconfigured)).toBe(false); + expect(getCustomerResourceCounts(unconfigured)).toEqual(noMcpCounts); + }); + + it.each([ + { mcp_servers: ["server-id"] }, + { mcp_access_groups: ["group-id"] }, + { mcp_toolsets: ["toolset-id"] }, + { mcp_tool_permissions: { "server-id": ["read"] } }, + ])("recognizes a grant from each MCP source (%j)", (grant) => { + expect(customerHasMcpAccess({ ...customer, object_permission: { ...emptyPermission, ...grant } })).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.ts new file mode 100644 index 00000000000..dbf30f0c7ac --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.ts @@ -0,0 +1,34 @@ +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; + +export type CustomerResourceCounts = { + models: number | "all"; + servers: number; + accessGroups: number; + toolsets: number; +}; + +export function getCustomerResourceCounts(customer: EndUser): CustomerResourceCounts { + const models = customer.models ?? []; + const permission = customer.object_permission; + const servers = permission?.mcp_servers ?? []; + const modelCount = + models.length === 0 || models.includes("all-proxy-models") || models.includes("*") ? "all" : new Set(models).size; + if (servers.includes(NO_MCP_SERVERS_SENTINEL)) { + return { models: modelCount, servers: 0, accessGroups: 0, toolsets: 0 }; + } + const serverIdentifiers = [...servers, ...Object.keys(permission?.mcp_tool_permissions ?? {})].filter( + (server) => server !== ALL_PROXY_MCP_SERVERS_SENTINEL && server !== NO_MCP_SERVERS_SENTINEL, + ); + return { + models: modelCount, + servers: new Set(serverIdentifiers).size, + accessGroups: new Set(permission?.mcp_access_groups ?? []).size, + toolsets: new Set(permission?.mcp_toolsets ?? []).size, + }; +} + +export function customerHasMcpAccess(customer: EndUser): boolean { + const counts = getCustomerResourceCounts(customer); + return counts.servers > 0 || counts.accessGroups > 0 || counts.toolsets > 0; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts index 057818a0ed3..46b17dd0c61 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts @@ -1,13 +1,21 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { apiClient } from "@/components/networking"; -import { createCustomer, deleteCustomers, setCustomerBlockedState, updateCustomer } from "./customerApi"; +import { + createCustomer, + deleteCustomers, + setCustomerBlockedState, + setCustomersBlockedState, + updateCustomer, +} from "./customerApi"; vi.mock("@/components/networking", () => ({ apiClient: { post: vi.fn() } })); const post = vi.mocked(apiClient.post); describe("customer API", () => { - beforeEach(() => post.mockReset()); + beforeEach(() => { + post.mockReset(); + }); it.each([ { name: "create", action: createCustomer, path: "/customer/new" }, @@ -18,27 +26,92 @@ describe("customer API", () => { post.mockResolvedValue(response); await expect(action("access-token", body)).resolves.toEqual(response); - expect(post).toHaveBeenCalledWith(path, { accessToken: "access-token", body }); + expect(post).toHaveBeenCalledWith(path, { + accessToken: "access-token", + body, + }); }); it("posts the customer IDs to delete", async () => { const userIds = ["customer-id", "another-id"]; - const response = { deleted_customers: userIds.length, message: "Customers deleted" }; + const response = { + deleted_customers: userIds.length, + message: "Customers deleted", + }; post.mockResolvedValue(response); await expect(deleteCustomers("access-token", userIds)).resolves.toEqual(response); - expect(post).toHaveBeenCalledWith("/customer/delete", { accessToken: "access-token", body: { user_ids: userIds } }); + expect(post).toHaveBeenCalledWith("/customer/delete", { + accessToken: "access-token", + body: { user_ids: userIds }, + }); }); - it.each([true, false])("updates the customer's blocked flag to %s", async (blocked) => { + it.each([ + { + blocked: true, + path: "/customer/block", + response: { blocked_users: [{ user_id: "customer-id", blocked: true }] }, + }, + { + blocked: false, + path: "/customer/unblock", + response: { blocked_users: ["remaining-id"] }, + }, + ])("posts one customer ID to $path", async ({ blocked, path, response }) => { const userId = "customer-id"; - const response = { user_id: userId, blocked, spend: 0 }; post.mockResolvedValue(response); await expect(setCustomerBlockedState("access-token", { userId, blocked })).resolves.toEqual(response); - expect(post).toHaveBeenCalledWith("/customer/update", { + expect(post).toHaveBeenCalledExactlyOnceWith(path, { accessToken: "access-token", - body: { user_id: userId, blocked }, + body: { user_ids: [userId] }, }); }); + + it.each([true, false])("propagates per-row action failures when blocked is %s", async (blocked) => { + const error = new Error(blocked ? "Customer action failed" : "Blocked user check was never set"); + post.mockRejectedValue(error); + + await expect( + setCustomerBlockedState("access-token", { + userId: "customer-id", + blocked, + }), + ).rejects.toBe(error); + }); + + it.each([ + { + blocked: true, + path: "/customer/block", + response: { blocked_users: [{ user_id: "customer-id", blocked: true }] }, + }, + { + blocked: false, + path: "/customer/unblock", + response: { blocked_users: ["remaining-id"] }, + }, + ])("posts all selected customer IDs to $path", async ({ blocked, path, response }) => { + const userIds = ["customer-id", "another-id"]; + post.mockResolvedValue(response); + + await expect(setCustomersBlockedState("access-token", { userIds, blocked })).resolves.toEqual(response); + expect(post).toHaveBeenCalledExactlyOnceWith(path, { + accessToken: "access-token", + body: { user_ids: userIds }, + }); + }); + + it.each([true, false])("propagates bulk action failures when blocked is %s", async (blocked) => { + const error = new Error(blocked ? "Customer action failed" : "Blocked user check was never set"); + post.mockRejectedValue(error); + + await expect( + setCustomersBlockedState("access-token", { + userIds: ["customer-id", "another-id"], + blocked, + }), + ).rejects.toBe(error); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts index 731c9b2dfa6..b92c8418dc2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts @@ -4,9 +4,15 @@ import type { EndUser } from "./useCustomers"; type Schemas = components["schemas"]; -export type CreateCustomerRequest = Partial & { user_id: string }; +export type CreateCustomerRequest = Partial & { + user_id: string; +}; export type UpdateCustomerRequest = Partial & { user_id: string }; export type CustomerBlockedStateInput = { userId: string; blocked: boolean }; +export type CustomersBlockedStateInput = { + userIds: string[]; + blocked: boolean; +}; export const createCustomer = (accessToken: string, body: CreateCustomerRequest): Promise => apiClient.post("/customer/new", { accessToken, body }); @@ -20,7 +26,20 @@ export const deleteCustomers = (accessToken: string, userIds: string[]): Promise body: { user_ids: userIds }, }); +export const setCustomersBlockedState = ( + accessToken: string, + { userIds, blocked }: CustomersBlockedStateInput, +): Promise => + apiClient.post( + blocked ? "/customer/block" : "/customer/unblock", + { + accessToken, + body: { user_ids: userIds }, + }, + ); + export const setCustomerBlockedState = ( accessToken: string, { userId, blocked }: CustomerBlockedStateInput, -): Promise => updateCustomer(accessToken, { user_id: userId, blocked }); +): Promise => + setCustomersBlockedState(accessToken, { userIds: [userId], blocked }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx index f6500fd088a..0a5d522a417 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx @@ -9,11 +9,14 @@ import { useCustomer, useDeleteCustomers, useSetCustomerBlockedState, + useSetCustomersBlockedState, useUpdateCustomer, } from "./useCustomerMutations"; const authorized = vi.hoisted(() => vi.fn()); -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: authorized })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: authorized, +})); vi.mock("@/components/networking", () => ({ apiClient: { post: vi.fn() } })); vi.mock("@/lib/http/api", () => ({ $api: { useQuery: vi.fn() } })); @@ -21,7 +24,9 @@ const post = vi.mocked(apiClient.post); const query = vi.mocked($api.useQuery); const createWrapper = () => { - const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false }, mutations: { retry: false } } }); + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); const wrapper = ({ children }: { children: ReactNode }) => ( {children} ); @@ -30,7 +35,10 @@ const createWrapper = () => { beforeEach(() => { vi.clearAllMocks(); - authorized.mockReturnValue({ accessToken: "access-token", userRole: "Admin" }); + authorized.mockReturnValue({ + accessToken: "access-token", + userRole: "Admin", + }); }); describe("useCustomer", () => { @@ -45,8 +53,18 @@ describe("useCustomer", () => { }); it.each([ - { role: "Admin Viewer", accessToken: "access-token", id: "customer-id", enabled: true }, - { role: "Internal User", accessToken: "access-token", id: "customer-id", enabled: false }, + { + role: "Admin Viewer", + accessToken: "access-token", + id: "customer-id", + enabled: true, + }, + { + role: "Internal User", + accessToken: "access-token", + id: "customer-id", + enabled: false, + }, { role: "Admin", accessToken: null, id: "customer-id", enabled: false }, { role: "Admin", accessToken: "access-token", id: null, enabled: false }, ])("gates the query for $role with customer $id", ({ role, accessToken, id, enabled }) => { @@ -76,7 +94,10 @@ describe("customer mutations", () => { const { result } = renderHook(() => useUpdateCustomer(), { wrapper }); await act(async () => { - await result.current.mutateAsync({ user_id: "customer-id", alias: "Alias" }); + await result.current.mutateAsync({ + user_id: "customer-id", + alias: "Alias", + }); }); await waitFor(() => expect(result.current.isSuccess).toBe(true)); @@ -123,4 +144,126 @@ describe("customer mutations", () => { }); expect(post).not.toHaveBeenCalled(); }); + + it.each([ + { + blocked: true, + path: "/customer/block", + response: { blocked_users: [{ user_id: "customer-id", blocked: true }] }, + }, + { + blocked: false, + path: "/customer/unblock", + response: { blocked_users: ["remaining-id"] }, + }, + ])("invalidates customer data after a per-row request to $path", async ({ blocked, path, response }) => { + const { queryClient, wrapper } = createWrapper(); + const listKey = ["get", "/customer/list", {}]; + const detailKey = ["get", "/customer/info", { params: { query: { end_user_id: "customer-id" } } }]; + const unrelatedKey = ["get", "/key/list", {}]; + queryClient.setQueryData(listKey, []); + queryClient.setQueryData(detailKey, {}); + queryClient.setQueryData(unrelatedKey, []); + post.mockResolvedValue(response); + const { result } = renderHook(useSetCustomerBlockedState, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ userId: "customer-id", blocked })).resolves.toEqual(response); + }); + + expect(post).toHaveBeenCalledExactlyOnceWith(path, { + accessToken: "access-token", + body: { user_ids: ["customer-id"] }, + }); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(detailKey)?.isInvalidated).toBe(true); + expect(queryClient.getQueryState(unrelatedKey)?.isInvalidated).toBe(false); + }); + + it("propagates per-row unblock errors and retains customer cache data", async () => { + const { queryClient, wrapper } = createWrapper(); + const listKey = ["get", "/customer/list", {}]; + const detailKey = ["get", "/customer/info", { params: { query: { end_user_id: "customer-id" } } }]; + const customer = { user_id: "customer-id", blocked: true }; + queryClient.setQueryData(listKey, [customer]); + queryClient.setQueryData(detailKey, customer); + const error = new Error("Blocked user check was never set"); + post.mockRejectedValue(error); + const { result } = renderHook(useSetCustomerBlockedState, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ userId: "customer-id", blocked: false })).rejects.toBe(error); + }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(queryClient.getQueryData(listKey)).toEqual([customer]); + expect(queryClient.getQueryData(detailKey)).toEqual(customer); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(false); + expect(queryClient.getQueryState(detailKey)?.isInvalidated).toBe(false); + }); + + it.each([true, false])("invalidates customer data after bulk blocked state becomes %s", async (blocked) => { + const { queryClient, wrapper } = createWrapper(); + const userIds = ["customer-id", "another-id"]; + const listKey = ["get", "/customer/list", {}]; + const detailKeys = userIds.map((id) => ["get", "/customer/info", { params: { query: { end_user_id: id } } }]); + queryClient.setQueryData(listKey, []); + detailKeys.forEach((key) => queryClient.setQueryData(key, {})); + const response = blocked + ? { + blocked_users: userIds.map((user_id) => ({ + user_id, + blocked: true, + })), + } + : { blocked_users: ["remaining-id"] }; + post.mockResolvedValue(response); + const { result } = renderHook(useSetCustomersBlockedState, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ userIds, blocked })).resolves.toEqual(response); + }); + + expect(post).toHaveBeenCalledExactlyOnceWith(blocked ? "/customer/block" : "/customer/unblock", { + accessToken: "access-token", + body: { user_ids: userIds }, + }); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(true); + detailKeys.forEach((key) => expect(queryClient.getQueryState(key)?.isInvalidated).toBe(true)); + }); + + it.each([true, false])("retains customer cache data when a bulk action fails with blocked %s", async (blocked) => { + const { queryClient, wrapper } = createWrapper(); + const listKey = ["get", "/customer/list", {}]; + const detailKey = ["get", "/customer/info", { params: { query: { end_user_id: "customer-id" } } }]; + const customer = { user_id: "customer-id", blocked: !blocked }; + queryClient.setQueryData(listKey, [customer]); + queryClient.setQueryData(detailKey, customer); + const error = new Error(blocked ? "Bulk block failed" : "Blocked user check was never set"); + post.mockRejectedValue(error); + const { result } = renderHook(useSetCustomersBlockedState, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ userIds: ["customer-id"], blocked })).rejects.toBe(error); + }); + + await waitFor(() => expect(result.current.isError).toBe(true)); + expect(queryClient.getQueryData(listKey)).toEqual([customer]); + expect(queryClient.getQueryData(detailKey)).toEqual(customer); + expect(queryClient.getQueryState(listKey)?.isInvalidated).toBe(false); + expect(queryClient.getQueryState(detailKey)?.isInvalidated).toBe(false); + }); + + it("rejects bulk blocked state changes without an access token", async () => { + authorized.mockReturnValue({ accessToken: null, userRole: "Admin" }); + const { wrapper } = createWrapper(); + const { result } = renderHook(useSetCustomersBlockedState, { wrapper }); + + await act(async () => { + await expect(result.current.mutateAsync({ userIds: ["customer-id"], blocked: true })).rejects.toThrow( + "Access token is required", + ); + }); + expect(post).not.toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts index 51c1cd8d9b0..10aebda76b8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.ts @@ -2,7 +2,13 @@ import { useMutation, useQueryClient } from "@tanstack/react-query"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { $api } from "@/lib/http/api"; import { all_admin_roles } from "@/utils/roles"; -import { createCustomer, deleteCustomers, setCustomerBlockedState, updateCustomer } from "./customerApi"; +import { + createCustomer, + deleteCustomers, + setCustomerBlockedState, + setCustomersBlockedState, + updateCustomer, +} from "./customerApi"; export const useCustomer = (id: string | null) => { const { accessToken, userRole } = useAuthorized(); @@ -34,3 +40,4 @@ export const useCreateCustomer = () => useCustomerMutation(createCustomer); export const useUpdateCustomer = () => useCustomerMutation(updateCustomer); export const useDeleteCustomers = () => useCustomerMutation(deleteCustomers); export const useSetCustomerBlockedState = () => useCustomerMutation(setCustomerBlockedState); +export const useSetCustomersBlockedState = () => useCustomerMutation(setCustomersBlockedState); From 4e4fdae0b73dbd89269205b25d0c1bf1e9f31035 Mon Sep 17 00:00:00 2001 From: lewis Date: Mon, 12 Oct 2026 17:46:24 -0400 Subject: [PATCH 11/56] feat(proxy): add customer bulk updates --- litellm/proxy/_types.py | 8 + litellm/proxy/auth/route_checks.py | 4 + .../customer_endpoints.py | 215 ++++++---- .../customer_endpoints.py | 20 +- .../test_customer_endpoints.py | 394 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 173 +++++++- 6 files changed, 737 insertions(+), 77 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b6c81f17761..c2193c19ef9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1834,6 +1834,14 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): object_permission: LiteLLM_ObjectPermissionBase | None = None +class BulkUpdateCustomerRequest(LiteLLMPydanticObjectBase): + model_config = ConfigDict(frozen=True) + + customers: tuple[UpdateCustomerRequest, ...] = Field( + min_length=1, max_length=500, description="Customer updates processed independently in request order" + ) + + class DeleteCustomerRequest(LiteLLMPydanticObjectBase): """ Delete multiple Customers diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index ae129ca47b5..d668d750d34 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -33,6 +33,8 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset( "/user/delete", "/management/v1/users/bulk_delete", "/user/bulk_update", + "/customer/bulk_update", + "/end_user/bulk_update", # team "/team/new", "/management/v1/teams/{team_id}/members/bulk_delete", @@ -923,6 +925,8 @@ class RouteChecks: "/user/delete", "/management/v1/users/bulk_delete", "/user/bulk_update", + "/customer/bulk_update", + "/end_user/bulk_update", "/team/new", "/management/v1/teams/{team_id}/members/bulk_delete", "/management/v1/teams/{team_id}/members/bulk_update", diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index d604c58d1f5..b1ff0f35b69 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -14,7 +14,7 @@ from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timedelta from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypeVar, overload +from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, Protocol, TypeVar, overload import fastapi from fastapi import APIRouter, Depends, HTTPException, Request @@ -51,7 +51,9 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( ) from litellm.types.proxy.management_endpoints.customer_endpoints import ( BlockUsersResponse, + BulkUpdateCustomerResponse, CustomerResponse, + CustomerUpdateResult, DeleteCustomersResponse, UnblockUsersResponse, ) @@ -593,76 +595,10 @@ async def end_user_info( raise handle_exception_on_proxy(e) -@router.post( - "/customer/update", - tags=["Customer Management"], - dependencies=[Depends(user_api_key_auth)], - response_model=CustomerResponse, -) -@router.post( - "/end_user/update", - tags=["Customer Management"], - include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], -) -async def update_end_user( +async def _update_single_customer( data: UpdateCustomerRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> CustomerResponse: - """ - Example curl - - Parameters: - - user_id: str - - alias: Optional[str] = None # human-friendly alias - - blocked: bool = False # allow/disallow requests for this end-user - - max_budget: Optional[float] = None - - budget_id: Optional[str] = None # give either a budget_id or max_budget - - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) - - models: Optional[list[str]] = None # omitted or null leaves the allowlist unchanged; an empty list clears it - - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permissions to control access to resources. - Supported fields: - * mcp_servers: List[str] - List of allowed MCP server IDs - * mcp_access_groups: List[str] - List of MCP access group names - * mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names - * vector_stores: List[str] - List of allowed vector store IDs - * agents: List[str] - List of allowed agent IDs - * agent_access_groups: List[str] - List of agent access group names - Example: {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]} - IF null or {} then no object-level restrictions apply. - - Example curl: - ``` - curl --location 'http://0.0.0.0:4000/customer/update' \ - --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ - --header 'Content-Type: application/json' \ - --data '{ - "user_id": "test-litellm-user-4", - "budget_id": "paid_tier", - "models": ["gpt-4o-mini"] - }' - - # Updating object permissions - curl -L -X POST 'http://localhost:4000/customer/update' \ - --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ - --header 'Content-Type: application/json' \ - --data '{ - "user_id": "user_1", - "object_permission": { - "mcp_servers": ["server_3"], - "vector_stores": ["vector_store_2", "vector_store_3"] - } - }' - - See below for all params - ``` - """ - + user_api_key_dict: UserAPIKeyAuth, +) -> CustomerResponse | ProxyException: from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client try: @@ -770,7 +706,144 @@ async def update_end_user( except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_end_user(): Exception occured - %s", e) - raise handle_exception_on_proxy(e) + return handle_exception_on_proxy(e) + + +@router.post( + "/customer/update", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + response_model=CustomerResponse, +) +@router.post( + "/end_user/update", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], +) +async def update_end_user( + data: UpdateCustomerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +) -> CustomerResponse: + """ + Example curl + + Parameters: + - user_id: str + - alias: Optional[str] = None # human-friendly alias + - blocked: bool = False # allow/disallow requests for this end-user + - max_budget: Optional[float] = None + - budget_id: Optional[str] = None # give either a budget_id or max_budget + - allowed_model_region: Optional[AllowedModelRegion] = ( + None # require all user requests to use models in this specific region + ) + - default_model: Optional[str] = ( + None # if no equivalent model in allowed region - default all requests to this model + ) + - models: Optional[list[str]] = None # omitted or null leaves the allowlist unchanged; an empty list clears it + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permissions to control access to resources. + Supported fields: + * mcp_servers: List[str] - List of allowed MCP server IDs + * mcp_access_groups: List[str] - List of MCP access group names + * mcp_toolsets: List[str] - List of allowed MCP toolset IDs + * mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names + * vector_stores: List[str] - List of allowed vector store IDs + * agents: List[str] - List of allowed agent IDs + * agent_access_groups: List[str] - List of agent access group names + Example: {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]} + Omitted permission fields retain their existing values. Supplying null or {} leaves permissions unchanged. + + Example curl: + ``` + curl --location 'http://0.0.0.0:4000/customer/update' \ + --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ + --header 'Content-Type: application/json' \ + --data '{ + "user_id": "test-litellm-user-4", + "budget_id": "paid_tier", + "models": ["gpt-4o-mini"] + }' + + # Updating object permissions + curl -L -X POST 'http://localhost:4000/customer/update' \ + --header "Authorization: Bearer $LITELLM_MASTER_KEY" \ + --header 'Content-Type: application/json' \ + --data '{ + "user_id": "user_1", + "object_permission": { + "mcp_servers": ["server_3"], + "vector_stores": ["vector_store_2", "vector_store_3"] + } + }' + + See below for all params + ``` + """ + + result: Final = await _update_single_customer(data=data, user_api_key_dict=user_api_key_dict) + if isinstance(result, ProxyException): + raise result + return result + + +async def _update_customer_result( + data: UpdateCustomerRequest, + user_api_key_dict: UserAPIKeyAuth, +) -> CustomerUpdateResult: + result: Final = await _update_single_customer(data=data, user_api_key_dict=user_api_key_dict) + if isinstance(result, ProxyException): + return CustomerUpdateResult(user_id=data.user_id, success=False, error=str(result)) + return CustomerUpdateResult(user_id=data.user_id, success=True, updated_customer=result) + + +@router.post( + "/customer/bulk_update", + tags=["Customer Management"], + dependencies=[Depends(user_api_key_auth)], + response_model=BulkUpdateCustomerResponse, +) +@router.post( + "/end_user/bulk_update", + tags=["Customer Management"], + include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], + response_model=BulkUpdateCustomerResponse, +) +async def bulk_customer_update( + data: BulkUpdateCustomerRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> BulkUpdateCustomerResponse: + """Update up to 500 customers independently using the same fields as `/customer/update`. + + Supply a `customers` array of customer updates. Each result contains the `user_id`, + `success`, and either `updated_customer` or `error`, in request order. A failed update + does not prevent the remaining customers from being updated. Object permission fields + omitted from an update retain their existing values, and supplied permission lists + replace that field. Only proxy admins can update customers. + + Example request body: + ```json + { + "customers": [ + {"user_id": "customer-1", "object_permission": {"mcp_toolsets": ["toolset-1"]}}, + {"user_id": "customer-2", "models": ["model-1"]} + ] + } + ``` + """ + results: Final = tuple( + [ + await _update_customer_result(data=customer, user_api_key_dict=user_api_key_dict) + for customer in data.customers + ] + ) + successful_updates: Final = sum(result.success for result in results) + return BulkUpdateCustomerResponse( + results=results, + total_requested=len(results), + successful_updates=successful_updates, + failed_updates=len(results) - successful_updates, + ) @router.post( diff --git a/litellm/types/proxy/management_endpoints/customer_endpoints.py b/litellm/types/proxy/management_endpoints/customer_endpoints.py index c1c4c496ede..4a4e778f4ec 100644 --- a/litellm/types/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/types/proxy/management_endpoints/customer_endpoints.py @@ -1,4 +1,4 @@ -from pydantic import Field +from pydantic import ConfigDict, Field from litellm.models.budget import LiteLLM_BudgetTableFull from litellm.models.end_user import LiteLLM_EndUserTable @@ -27,3 +27,21 @@ class UnblockUsersResponse(LiteLLMBaseModel): class DeleteCustomersResponse(LiteLLMBaseModel): deleted_customers: int message: str + + +class CustomerUpdateResult(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + user_id: str + success: bool + updated_customer: CustomerResponse | None = None + error: str | None = None + + +class BulkUpdateCustomerResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + results: tuple[CustomerUpdateResult, ...] + total_requested: int + successful_updates: int + failed_updates: int diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index 2ac7b68766b..ec2c06caf36 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from contextlib import contextmanager +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -8,14 +9,22 @@ 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 TypeAdapter from litellm.proxy._types import ( LiteLLM_EndUserTable, LiteLLM_ObjectPermissionTable, + LiteLLM_UserTable, LitellmUserRoles, ProxyException, ) +from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + end_user_cache_key, + end_user_restricted_registry_cache_key, +) from litellm.proxy.management_endpoints.customer_endpoints import _should_update_field, router from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, @@ -918,9 +927,7 @@ def test_customer_new_forwards_models_to_db(mock_prisma_client, mock_user_api_ke @pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) -def test_customer_new_rejects_a_duration_that_never_advances( - mock_prisma_client, mock_user_api_key_auth, bad_duration -): +def test_customer_new_rejects_a_duration_that_never_advances(mock_prisma_client, mock_user_api_key_auth, bad_duration): """A zero-length window resets to "now", leaving the customer's budget row permanently due for the reset job to re-read every tick.""" mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) @@ -1189,3 +1196,384 @@ def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_cl "end_user_id:c2", "end_user_restricted_registry", ] + + +class _CustomerBatchDatabase: + def __init__( + self, + customers: tuple[LiteLLM_EndUserTable, ...], + permissions: tuple[LiteLLM_ObjectPermissionTable, ...] = (), + ) -> None: + self.customers: Mapping[str, LiteLLM_EndUserTable] = MappingProxyType( + {customer.user_id: customer for customer in customers} + ) + self.permissions: Mapping[str, LiteLLM_ObjectPermissionTable] = MappingProxyType( + {permission.object_permission_id: permission for permission in permissions} + ) + + async def find_customer( + self, *, where: Mapping[str, str], include: Mapping[str, bool] + ) -> LiteLLM_EndUserTable | None: + return self.customers.get(where["user_id"]) + + async def update_customer( + self, *, where: Mapping[str, str], data: Mapping[str, object], include: Mapping[str, bool] + ) -> LiteLLM_EndUserTable: + existing: Final = self.customers[where["user_id"]] + updated: Final = LiteLLM_EndUserTable.model_validate({**existing.model_dump(), **data}) + permission: Final = self.permissions.get(updated.object_permission_id or "") + saved: Final = updated.model_copy(update={"object_permission": permission}) + self.customers = MappingProxyType({**self.customers, saved.user_id: saved}) + return saved + + async def find_permission(self, *, where: Mapping[str, str]) -> LiteLLM_ObjectPermissionTable | None: + return self.permissions.get(where["object_permission_id"]) + + async def upsert_permission( + self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> LiteLLM_ObjectPermissionTable: + stored: Final = data["update"] if where["object_permission_id"] in self.permissions else data["create"] + tool_permissions: Final = stored.get("mcp_tool_permissions") + decoded: Final = ( + TypeAdapter(dict[str, list[str]]).validate_json(tool_permissions) + if isinstance(tool_permissions, str) + else tool_permissions + ) + permission: Final = LiteLLM_ObjectPermissionTable.model_validate({**stored, "mcp_tool_permissions": decoded}) + self.permissions = MappingProxyType({**self.permissions, permission.object_permission_id: permission}) + return permission + + def attach(self, prisma_client: MagicMock) -> None: + prisma_client.configure_mock( + **{ + "db.litellm_endusertable.find_first": AsyncMock(side_effect=self.find_customer), + "db.litellm_endusertable.update": AsyncMock(side_effect=self.update_customer), + "db.litellm_objectpermissiontable.find_unique": AsyncMock(side_effect=self.find_permission), + "db.litellm_objectpermissiontable.upsert": AsyncMock(side_effect=self.upsert_permission), + } + ) + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +def test_customer_bulk_update_applies_each_customers_fields( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + database: Final = _CustomerBatchDatabase( + ( + LiteLLM_EndUserTable(user_id="bulk-a", alias="Original A", blocked=True, models=["model-a"]), + LiteLLM_EndUserTable(user_id="bulk-b", alias="Original B", blocked=False, models=["model-b"]), + ) + ) + database.attach(mock_prisma_client) + + response: Final = client.post( + route, + json={ + "customers": [ + {"user_id": "bulk-a", "alias": "Updated A", "models": ["replacement-model"]}, + {"user_id": "bulk-b", "blocked": True, "models": []}, + ] + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, response.text + assert TypeAdapter(dict[str, object]).validate_json(response.content) == { + "total_requested": 2, + "successful_updates": 2, + "failed_updates": 0, + "results": [ + { + "user_id": "bulk-a", + "success": True, + "updated_customer": CustomerResponse( + user_id="bulk-a", alias="Updated A", blocked=True, models=["replacement-model"] + ).model_dump(mode="json"), + "error": None, + }, + { + "user_id": "bulk-b", + "success": True, + "updated_customer": CustomerResponse(user_id="bulk-b", alias="Original B", blocked=True).model_dump( + mode="json" + ), + "error": None, + }, + ], + } + assert database.customers["bulk-a"].alias == "Updated A" + assert database.customers["bulk-a"].models == ["replacement-model"] + assert database.customers["bulk-b"].models == [] + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +def test_customer_bulk_update_reports_unknown_middle_item_and_updates_later_customer( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + database: Final = _CustomerBatchDatabase( + ( + LiteLLM_EndUserTable(user_id="bulk-a", blocked=False), + LiteLLM_EndUserTable(user_id="bulk-b", blocked=False), + ) + ) + database.attach(mock_prisma_client) + + response: Final = client.post( + route, + json={ + "customers": [ + {"user_id": "bulk-a", "alias": "Updated A"}, + {"user_id": "bulk-missing", "alias": "Unknown"}, + {"user_id": "bulk-b", "alias": "Updated B"}, + ] + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, response.text + assert TypeAdapter(dict[str, object]).validate_json(response.content) == { + "total_requested": 3, + "successful_updates": 2, + "failed_updates": 1, + "results": [ + { + "user_id": "bulk-a", + "success": True, + "updated_customer": CustomerResponse(user_id="bulk-a", alias="Updated A", blocked=False).model_dump( + mode="json" + ), + "error": None, + }, + { + "user_id": "bulk-missing", + "success": False, + "updated_customer": None, + "error": "End User Id=bulk-missing does not exist in db", + }, + { + "user_id": "bulk-b", + "success": True, + "updated_customer": CustomerResponse(user_id="bulk-b", alias="Updated B", blocked=False).model_dump( + mode="json" + ), + "error": None, + }, + ], + } + assert database.customers["bulk-a"].alias == "Updated A" + assert database.customers["bulk-b"].alias == "Updated B" + assert "bulk-missing" not in database.customers + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +def test_customer_bulk_update_toolsets_only_preserves_each_customers_other_permissions( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + permissions: Final = tuple( + LiteLLM_ObjectPermissionTable( + object_permission_id=f"permission-{suffix}", + mcp_servers=[f"server-{suffix}"], + mcp_access_groups=[f"group-{suffix}"], + mcp_tool_permissions={f"server-{suffix}": [f"read-{suffix}"]}, + vector_stores=[f"vector-store-{suffix}"], + agents=[f"agent-{suffix}"], + agent_access_groups=[f"agent-group-{suffix}"], + models=[f"permission-model-{suffix}"], + mcp_toolsets=[f"old-toolset-{suffix}"], + blocked_tools=[f"blocked-tool-{suffix}"], + search_tools=[f"search-{suffix}"], + mcp_tool_search_enabled=suffix == "a", + skills=[f"skill-{suffix}"], + ) + for suffix in ("a", "b") + ) + database: Final = _CustomerBatchDatabase( + tuple( + LiteLLM_EndUserTable(user_id=f"bulk-{suffix}", blocked=False, object_permission_id=f"permission-{suffix}") + for suffix in ("a", "b") + ), + permissions, + ) + database.attach(mock_prisma_client) + + response: Final = client.post( + route, + json={ + "customers": [ + {"user_id": "bulk-a", "object_permission": {"mcp_toolsets": ["new-toolset"]}}, + {"user_id": "bulk-b", "object_permission": {"mcp_toolsets": ["new-toolset"]}}, + ] + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, response.text + body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) + assert (body["total_requested"], body["successful_updates"], body["failed_updates"]) == (2, 2, 0) + results: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(body["results"]) + assert len(results) == 2 + for original, result in zip(permissions, results, strict=True): + expected: Final = original.model_copy(update={"mcp_toolsets": ["new-toolset"]}) + customer: Final = CustomerResponse.model_validate(result["updated_customer"]) + assert result["success"] is True + assert result["error"] is None + assert customer.object_permission == expected + assert database.permissions[original.object_permission_id] == expected + assert database.customers[customer.user_id].object_permission == expected + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +@pytest.mark.parametrize("role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)) +def test_customer_bulk_update_rejects_non_admins_without_writing( + mock_prisma_client: MagicMock, route: str, role: LitellmUserRoles +) -> None: + original: Final = LiteLLM_EndUserTable(user_id="bulk-a", alias="Original", blocked=False) + database: Final = _CustomerBatchDatabase((original,)) + database.attach(mock_prisma_client) + payload: Final[dict[str, object]] = {"customers": [{"user_id": "bulk-a", "alias": "Updated"}]} + + def authorize_role(request: Request) -> UserAPIKeyAuth: + auth: Final = UserAPIKeyAuth(user_id="bulk-non-admin", user_role=role) + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=LiteLLM_UserTable(user_id="bulk-non-admin", user_role=role), + _user_role=role, + route=request.url.path, + request=request, + valid_token=auth, + request_data=payload, + ) + except HTTPException: + raise + except Exception as error: + raise HTTPException(status_code=403, detail=str(error)) from error + return auth + + with patch.dict(app.dependency_overrides, {user_api_key_auth: authorize_role}): + response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == 403, response.text + assert route in response.text + assert database.customers == {"bulk-a": original} + mock_prisma_client.db.litellm_endusertable.update.assert_not_awaited() + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +def test_customer_bulk_update_evicts_every_updated_id_and_registry_only( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + database: Final = _CustomerBatchDatabase( + ( + LiteLLM_EndUserTable(user_id="bulk-a", blocked=False), + LiteLLM_EndUserTable(user_id="bulk-b", blocked=False), + LiteLLM_EndUserTable(user_id="bulk-untouched", blocked=False), + ) + ) + database.attach(mock_prisma_client) + cache: Final = UserApiKeyCache() + for user_id in ("bulk-a", "bulk-b", "bulk-missing", "bulk-untouched"): + cache.set_cache(key=end_user_cache_key(user_id), value=f"cached-{user_id}") + cache.set_cache(key=end_user_restricted_registry_cache_key(), value="cached-registry") + + with patch.multiple("litellm.proxy.proxy_server", user_api_key_cache=cache, redis_usage_cache=None): + response: Final = client.post( + route, + json={ + "customers": [ + {"user_id": "bulk-a", "alias": "Updated A"}, + {"user_id": "bulk-missing", "alias": "Unknown"}, + {"user_id": "bulk-b", "alias": "Updated B"}, + ] + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, response.text + body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) + assert (body["successful_updates"], body["failed_updates"]) == (2, 1) + assert cache.get_cache(key=end_user_cache_key("bulk-a")) is None + assert cache.get_cache(key=end_user_cache_key("bulk-b")) is None + assert cache.get_cache(key=end_user_restricted_registry_cache_key()) is None + assert cache.get_cache(key=end_user_cache_key("bulk-missing")) == "cached-bulk-missing" + assert cache.get_cache(key=end_user_cache_key("bulk-untouched")) == "cached-bulk-untouched" + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +@pytest.mark.parametrize( + "customers", + ([], [{"user_id": "bulk-a"}] * 501, [{"user_id": "bulk-a"}, {"alias": "Missing ID"}]), + ids=("empty", "over-limit", "malformed-item"), +) +def test_customer_bulk_update_rejects_invalid_batches_before_writing( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str, customers: list[dict[str, object]] +) -> None: + original: Final = LiteLLM_EndUserTable(user_id="bulk-a", alias="Original", blocked=False) + database: Final = _CustomerBatchDatabase((original,)) + database.attach(mock_prisma_client) + + response: Final = client.post( + route, + json={"customers": customers}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 422, response.text + assert database.customers == {"bulk-a": original} + mock_prisma_client.db.litellm_endusertable.find_first.assert_not_awaited() + mock_prisma_client.db.litellm_endusertable.update.assert_not_awaited() + + +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +@pytest.mark.parametrize("role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.ORG_ADMIN)) +def test_customer_bulk_update_rejects_org_admin_context_without_writing( + mock_prisma_client: MagicMock, route: str, role: LitellmUserRoles +) -> None: + original: Final = LiteLLM_EndUserTable(user_id="bulk-a", alias="Original", blocked=False) + database: Final = _CustomerBatchDatabase((original,)) + database.attach(mock_prisma_client) + payload: Final[dict[str, object]] = { + "organization_id": "bulk-organization", + "customers": [{"user_id": "bulk-a", "alias": "Updated"}], + } + user: Final = LiteLLM_UserTable.model_validate( + { + "user_id": "bulk-org-admin", + "user_role": role, + "organization_memberships": [ + { + "user_id": "bulk-org-admin", + "organization_id": "bulk-organization", + "user_role": LitellmUserRoles.ORG_ADMIN.value, + "created_at": "2026-01-01T00:00:00+00:00", + "updated_at": "2026-01-01T00:00:00+00:00", + } + ], + } + ) + + def authorize_role(request: Request) -> UserAPIKeyAuth: + auth: Final = UserAPIKeyAuth(user_id=user.user_id, user_role=role) + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user, + _user_role=role, + route=request.url.path, + request=request, + valid_token=auth, + request_data=payload, + ) + except HTTPException: + raise + except Exception as error: + raise HTTPException(status_code=403, detail=str(error)) from error + return auth + + with patch.dict(app.dependency_overrides, {user_api_key_auth: authorize_role}): + response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == 403, response.text + assert route in response.text + assert database.customers == {"bulk-a": original} + mock_prisma_client.db.litellm_endusertable.update.assert_not_awaited() + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_awaited() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ea40e2d90b9..db1ba848924 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4169,6 +4169,42 @@ export interface paths { patch?: never; trace?: never; }; + "/customer/bulk_update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Bulk Customer Update + * @description Update up to 500 customers independently using the same fields as `/customer/update`. + * + * Supply a `customers` array of customer updates. Each result contains the `user_id`, + * `success`, and either `updated_customer` or `error`, in request order. A failed update + * does not prevent the remaining customers from being updated. Object permission fields + * omitted from an update retain their existing values, and supplied permission lists + * replace that field. Only proxy admins can update customers. + * + * Example request body: + * ```json + * { + * "customers": [ + * {"user_id": "customer-1", "object_permission": {"mcp_toolsets": ["toolset-1"]}}, + * {"user_id": "customer-2", "models": ["model-1"]} + * ] + * } + * ``` + */ + post: operations["bulk_customer_update_customer_bulk_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/customer/daily/activity": { parameters: { query?: never; @@ -4498,12 +4534,13 @@ export interface paths { * Supported fields: * * mcp_servers: List[str] - List of allowed MCP server IDs * * mcp_access_groups: List[str] - List of MCP access group names + * * mcp_toolsets: List[str] - List of allowed MCP toolset IDs * * mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names * * vector_stores: List[str] - List of allowed vector store IDs * * agents: List[str] - List of allowed agent IDs * * agent_access_groups: List[str] - List of agent access group names * Example: {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]} - * IF null or {} then no object-level restrictions apply. + * Omitted permission fields retain their existing values. Supplying null or {} leaves permissions unchanged. * * Example curl: * ``` @@ -4846,6 +4883,42 @@ export interface paths { patch?: never; trace?: never; }; + "/end_user/bulk_update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Bulk Customer Update + * @description Update up to 500 customers independently using the same fields as `/customer/update`. + * + * Supply a `customers` array of customer updates. Each result contains the `user_id`, + * `success`, and either `updated_customer` or `error`, in request order. A failed update + * does not prevent the remaining customers from being updated. Object permission fields + * omitted from an update retain their existing values, and supplied permission lists + * replace that field. Only proxy admins can update customers. + * + * Example request body: + * ```json + * { + * "customers": [ + * {"user_id": "customer-1", "object_permission": {"mcp_toolsets": ["toolset-1"]}}, + * {"user_id": "customer-2", "models": ["model-1"]} + * ] + * } + * ``` + */ + post: operations["bulk_customer_update_end_user_bulk_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/end_user/daily/activity": { parameters: { query?: never; @@ -5175,12 +5248,13 @@ export interface paths { * Supported fields: * * mcp_servers: List[str] - List of allowed MCP server IDs * * mcp_access_groups: List[str] - List of MCP access group names + * * mcp_toolsets: List[str] - List of allowed MCP toolset IDs * * mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names * * vector_stores: List[str] - List of allowed vector store IDs * * agents: List[str] - List of allowed agent IDs * * agent_access_groups: List[str] - List of agent access group names * Example: {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]} - * IF null or {} then no object-level restrictions apply. + * Omitted permission fields retain their existing values. Supplying null or {} leaves permissions unchanged. * * Example curl: * ``` @@ -27720,6 +27794,25 @@ export interface components { /** Data */ data: components["schemas"]["TeamMemberDeleteResult"][]; }; + /** BulkUpdateCustomerRequest */ + BulkUpdateCustomerRequest: { + /** + * Customers + * @description Customer updates processed independently in request order + */ + customers: components["schemas"]["UpdateCustomerRequest"][]; + }; + /** BulkUpdateCustomerResponse */ + BulkUpdateCustomerResponse: { + /** Failed Updates */ + failed_updates: number; + /** Results */ + results: components["schemas"]["CustomerUpdateResult"][]; + /** Successful Updates */ + successful_updates: number; + /** Total Requested */ + total_requested: number; + }; /** * BulkUpdateKeyRequest * @description Request for bulk key updates @@ -30902,6 +30995,16 @@ export interface components { /** User Id */ user_id: string; }; + /** CustomerUpdateResult */ + CustomerUpdateResult: { + /** Error */ + error?: string | null; + /** Success */ + success: boolean; + updated_customer?: components["schemas"]["CustomerResponse"] | null; + /** User Id */ + user_id: string; + }; /** * CyberArkConfig * @description Configuration for CyberArk Conjur secret manager integration. @@ -56348,6 +56451,39 @@ export interface operations { }; }; }; + bulk_customer_update_customer_bulk_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BulkUpdateCustomerRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["BulkUpdateCustomerResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_daily_activity_customer_daily_activity_get: { parameters: { query?: { @@ -57367,6 +57503,39 @@ export interface operations { }; }; }; + bulk_customer_update_end_user_bulk_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BulkUpdateCustomerRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["BulkUpdateCustomerResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_daily_activity_end_user_daily_activity_get: { parameters: { query?: { From 51e8c394a4b83e76750e57c6b20a4b77e2e74230 Mon Sep 17 00:00:00 2001 From: lewis Date: Tue, 13 Oct 2026 21:13:11 -0400 Subject: [PATCH 12/56] fix(permissions): preserve unsent grants in concurrent updates --- .../proxy/agent_endpoints/agent_registry.py | 2 +- .../key_management_endpoints.py | 2 +- .../organization_endpoints.py | 2 +- .../object_permission_utils.py | 37 +++++-- .../agent_endpoints/test_agent_registry.py | 8 +- .../test_customer_endpoints.py | 2 +- .../test_key_management_endpoints.py | 6 +- .../test_organization_endpoints.py | 8 +- .../test_object_permission_utils.py | 101 +++++++++++++++++- 9 files changed, 146 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 7929f67720d..c83ff141e88 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -162,7 +162,7 @@ async def _permission_write( if existing_id is None: created: Final[AgentPermissionWrite] = {"create": prepared.record} return created - updated: Final[AgentPermissionWrite] = {"update": prepared.record} + updated: Final[AgentPermissionWrite] = {"update": prepared.update_record} return updated diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 782a7e1db3c..7bad735f7f7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2760,7 +2760,7 @@ async def _write_prepared_key_update_object_permission( await ObjectPermissionRepository(prisma_client).table.upsert( where={"object_permission_id": upsert.object_permission_id}, - data={"create": upsert.record, "update": upsert.record}, + data={"create": upsert.record, "update": upsert.update_record}, ) return MappingProxyType({**data_json, "object_permission_id": upsert.object_permission_id}) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index b8ffec639c0..15ef49a7a3a 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -988,7 +988,7 @@ async def update_organization_v2( where={"object_permission_id": object_permission_upsert.object_permission_id}, data={ "create": object_permission_upsert.record, - "update": object_permission_upsert.record, + "update": object_permission_upsert.update_record, }, ) if budget_updates: diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 969ebf2cf8c..b061f264b47 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -24,6 +24,7 @@ from litellm.repositories.object_permission_repository import ObjectPermissionRe from litellm.repositories.table_repositories import MCPServerRepository if TYPE_CHECKING: + from prisma import Prisma from prisma import models as prisma_models from litellm.proxy._types import ( @@ -32,6 +33,11 @@ if TYPE_CHECKING: ) +@dataclass(frozen=True, slots=True) +class ObjectPermissionClient: + db: "Prisma" + + async def attach_object_permission_to_dict( data_dict: dict[str, object], prisma_client: PrismaClient, @@ -77,19 +83,21 @@ async def attach_object_permission_to_dict( class ObjectPermissionUpsert: object_permission_id: str record: dict[str, object] + update_record: Mapping[str, object] async def prepare_object_permission_upsert( new_object_permission: Mapping[str, object], existing_object_permission_id: str | None, - prisma_client: PrismaClient, + prisma_client: PrismaClient | ObjectPermissionClient, ) -> ObjectPermissionUpsert: """ Read-and-merge half of an object permission upsert; performs no writes. Merges the sent grants over the existing row (looked up by ``existing_object_permission_id``, or a fresh uuid when the entity has none) and - returns the id plus the full record to upsert. The id is pinned inside the record + returns the id, a merged create record, and an update record of only sent grants. + The id is pinned inside the create record because the column has ``@default(uuid())``, so a create without it would mint a different id than the one the caller links. ``mcp_tool_permissions`` is serialized to a JSON string to avoid GraphQL parsing issues (e.g. server IDs starting with @@ -113,10 +121,15 @@ async def prepare_object_permission_upsert( existing_mcp_tool_permissions=existing_fields.get("mcp_tool_permissions"), prisma_client=prisma_client, ) + requested_fields: Final = { + field: value + for field, value in new_object_permission.items() + if value is not None and field != "object_permission_id" + } # Omitted request fields become None; preserve stored grants and Prisma column defaults. merged: Final[dict[str, object]] = { **existing_fields, - **{field: value for field, value in new_object_permission.items() if value is not None}, + **requested_fields, "object_permission_id": object_permission_id, } record: Final[dict[str, object]] = { @@ -127,7 +140,15 @@ async def prepare_object_permission_upsert( else {} ), } - return ObjectPermissionUpsert(object_permission_id=object_permission_id, record=record) + update_record: Final = { + field: safe_dumps(value) if field == "mcp_tool_permissions" else value + for field, value in requested_fields.items() + } + return ObjectPermissionUpsert( + object_permission_id=object_permission_id, + record=record, + update_record=update_record, + ) async def handle_update_object_permission_common( @@ -175,7 +196,7 @@ async def handle_update_object_permission_common( where={"object_permission_id": upsert.object_permission_id}, data={ "create": upsert.record, - "update": upsert.record, + "update": upsert.update_record, }, ) @@ -263,7 +284,7 @@ def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: async def _get_db_mcp_servers_by_identifiers( identifiers: AbstractSet[str], - prisma_client: PrismaClient | None, + prisma_client: PrismaClient | ObjectPermissionClient | None, ) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]": if prisma_client is None or not identifiers: return [] @@ -282,7 +303,7 @@ async def _get_db_mcp_servers_by_identifiers( async def _resolve_mcp_server_identifiers_to_ids( identifiers: AbstractSet[str], - prisma_client: PrismaClient | None, + prisma_client: PrismaClient | ObjectPermissionClient | None, ) -> dict[str, set[str]]: """ Resolve MCP permission entries written as server_id, alias, or server_name @@ -339,7 +360,7 @@ def _mcp_tool_permission_entries(raw: object) -> Mapping[str, frozenset[str]]: async def reject_ambiguous_mcp_tool_permission_keys( new_mcp_tool_permissions: object, existing_mcp_tool_permissions: object, - prisma_client: PrismaClient | None, + prisma_client: PrismaClient | ObjectPermissionClient | None, ) -> None: """ A name or alias shared by several MCP servers cannot key ``mcp_tool_permissions``: diff --git a/tests/unit/proxy/agent_endpoints/test_agent_registry.py b/tests/unit/proxy/agent_endpoints/test_agent_registry.py index 7663f1d30e6..d2465213b55 100644 --- a/tests/unit/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/unit/proxy/agent_endpoints/test_agent_registry.py @@ -1485,10 +1485,10 @@ async def test_agent_permissions_are_written_atomically_with_the_registration(op permission: Final = write.call_args.kwargs["data"]["object_permission"][ "create" if operation == "create" else "update" ] - assert permission["models"] == ["new"] - if operation != "create": - assert permission["mcp_servers"] == ["slack"] - assert permission["object_permission_id"] == "permissions" + if operation == "create": + assert permission == {"object_permission_id": permission["object_permission_id"], "models": ["new"]} + else: + assert permission == {"models": ["new"]} client.db.litellm_objectpermissiontable.update.assert_not_called() client.db.litellm_objectpermissiontable.create.assert_not_called() diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index 2ac7b68766b..4f3e3147e56 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1019,7 +1019,7 @@ def test_customer_update_creates_partial_toolset_permission( "object_permission_id": where["object_permission_id"], "mcp_toolsets": ["toolset-id"], } - assert data["update"] == data["create"] + assert data["update"] == {"mcp_toolsets": ["toolset-id"]} return LiteLLM_ObjectPermissionTable.model_validate(data["create"]) mock_prisma_client.configure_mock( diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 71334e2321b..60950c4dce1 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -1135,7 +1135,7 @@ async def test_key_update_prepares_existing_object_permission_without_writing(): assert result == {"user_id": "user123", "object_permission_id": "existing_perm_id_123"} permission_upsert.assert_awaited_once_with( where={"object_permission_id": "existing_perm_id_123"}, - data={"create": upsert.record, "update": upsert.record}, + data={"create": upsert.record, "update": {"vector_stores": ["new_store"]}}, ) @@ -1179,7 +1179,7 @@ async def test_key_update_prepares_json_object_permission_and_upserts_new_row(): assert result["object_permission_id"] == upsert.object_permission_id permission_upsert.assert_awaited_once_with( where={"object_permission_id": upsert.object_permission_id}, - data={"create": upsert.record, "update": upsert.record}, + data={"create": upsert.record, "update": {"vector_stores": ["brand_new_store"]}}, ) @@ -1222,7 +1222,7 @@ async def test_key_update_recreates_missing_object_permission_with_existing_id() ) permission_upsert.assert_awaited_once_with( where={"object_permission_id": "missing_perm_id_789"}, - data={"create": upsert.record, "update": upsert.record}, + data={"create": upsert.record, "update": {"vector_stores": ["recreated_store"]}}, ) diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index a98f1344279..7ca72b7b2aa 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1056,8 +1056,12 @@ async def test_v2_wires_object_permission_onto_org_write(monkeypatch): upsert = prisma.tx.litellm_objectpermissiontable.upsert.await_args.kwargs assert upsert["where"] == {"object_permission_id": "op-123"} - assert upsert["data"]["update"]["mcp_servers"] == ["server-1"] - assert upsert["data"]["update"]["vector_stores"] == ["vs-1"] + assert upsert["data"]["create"] == { + "object_permission_id": "op-123", + "mcp_servers": ["server-1"], + "vector_stores": ["vs-1"], + } + assert upsert["data"]["update"] == {"vector_stores": ["vs-1"]} write_data = prisma.db.litellm_organizationtable.update.await_args.kwargs["data"] assert write_data["object_permission_id"] == "op-123" diff --git a/tests/unit/proxy/management_helpers/test_object_permission_utils.py b/tests/unit/proxy/management_helpers/test_object_permission_utils.py index 34201bde6e0..a5c2bfbb7e8 100644 --- a/tests/unit/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/unit/proxy/management_helpers/test_object_permission_utils.py @@ -1,5 +1,6 @@ +import asyncio import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -21,6 +22,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _resolve_team_allowed_mcp_servers, set_object_permission, enforce_all_proxy_mcp_servers_grant_is_admin_only, + handle_update_object_permission_common, prepare_object_permission_upsert, validate_key_mcp_servers_against_team, validate_key_search_tools_against_team, @@ -1590,6 +1592,103 @@ async def test_prepare_object_permission_upsert_preserves_unsent_grants_on_updat } +@pytest.mark.asyncio +async def test_permission_update_patch_keeps_false_and_empty_grants_without_changing_id() -> None: + original: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="permission-id", + mcp_servers=["existing-server"], + mcp_tool_permissions={"solo-id": ["read"]}, + mcp_toolsets=["existing-toolset"], + mcp_tool_search_enabled=True, + ) + prisma_client: Final = _make_ambiguity_prisma(existing_row_override=original) + + upsert: Final = await prepare_object_permission_upsert( + new_object_permission={ + "object_permission_id": "other-id", + "mcp_servers": None, + "mcp_tool_permissions": {}, + "mcp_toolsets": [], + "mcp_tool_search_enabled": False, + }, + existing_object_permission_id=original.object_permission_id, + prisma_client=prisma_client, + ) + + assert upsert.record == { + "object_permission_id": original.object_permission_id, + "mcp_servers": original.mcp_servers, + "mcp_tool_permissions": "{}", + "mcp_toolsets": [], + "mcp_tool_search_enabled": False, + } + assert upsert.update_record == { + "mcp_tool_permissions": "{}", + "mcp_toolsets": [], + "mcp_tool_search_enabled": False, + } + + +class _InterleavedPermissionTable: + def __init__(self, row: LiteLLM_ObjectPermissionTable) -> None: + self.row = row + self.read_count = 0 + self.reads_completed = asyncio.Event() + + async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_ObjectPermissionTable: + snapshot: Final = self.row + self.read_count += 1 + if self.read_count == 2: + self.reads_completed.set() + await self.reads_completed.wait() + return snapshot + + async def upsert( + self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> LiteLLM_ObjectPermissionTable: + self.row = LiteLLM_ObjectPermissionTable.model_validate({**self.row.model_dump(), **data["update"]}) + return self.row + + +@pytest.mark.asyncio +@pytest.mark.parametrize("toolsets_first", (False, True)) +async def test_partial_permission_updates_preserve_concurrently_changed_unsent_grants(toolsets_first: bool) -> None: + original: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="permission-id", + mcp_servers=["old-server"], + mcp_access_groups=["existing-group"], + mcp_toolsets=["old-toolset"], + mcp_tool_search_enabled=False, + skills=["existing-skill"], + ) + table: Final = _InterleavedPermissionTable(original) + prisma_client: Final = MagicMock() + prisma_client.configure_mock( + **{ + "db.litellm_objectpermissiontable.find_unique": AsyncMock(side_effect=table.find_unique), + "db.litellm_objectpermissiontable.upsert": AsyncMock(side_effect=table.upsert), + } + ) + server_patch: Final = {"mcp_servers": ["new-server"]} + toolset_patch: Final = {"mcp_toolsets": ["new-toolset"]} + patches: Final = (toolset_patch, server_patch) if toolsets_first else (server_patch, toolset_patch) + + await asyncio.gather( + *( + handle_update_object_permission_common( + data_json={"object_permission": permission}, + existing_object_permission_id=original.object_permission_id, + prisma_client=prisma_client, + ) + for permission in patches + ) + ) + + assert table.row == original.model_copy( + update={"mcp_servers": ["new-server"], "mcp_toolsets": ["new-toolset"]} + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "identifier, colliding_ids", From c25fe0d1221dd399f6a9216a5a3d13132c9c4f79 Mon Sep 17 00:00:00 2001 From: lewis Date: Tue, 13 Oct 2026 21:20:07 -0400 Subject: [PATCH 13/56] fix(customers): make bulk update items atomic --- .../customer_endpoints.py | 208 +++++----- .../test_customer_endpoints.py | 355 +++++++++++++++++- 2 files changed, 447 insertions(+), 116 deletions(-) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index b1ff0f35b69..f47a46ef38d 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -38,13 +38,16 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.proxy.management_endpoints.common_utils import validate_budget_duration from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports + ObjectPermissionClient, + ObjectPermissionUpsert, _set_object_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - handle_update_object_permission_common, + prepare_object_permission_upsert, set_object_permission, ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in +from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import EndUserRepository from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, @@ -61,6 +64,7 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import ( _RowT_co: Final = TypeVar("_RowT_co", covariant=True) _STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) _CLEARABLE_LIST_FIELDS: Final = frozenset({"models"}) +_CUSTOMER_UPDATE_LOCK_SQL: Final = 'SELECT "user_id" FROM "LiteLLM_EndUserTable" WHERE "user_id" = $1 FOR UPDATE' def _should_update_field(field: str, value: object, sent_fields: AbstractSet[str]) -> bool: @@ -291,34 +295,39 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: return None -async def _handle_customer_object_permission_update( - non_default_values: dict[str, object], - end_user_table_data_typed: LiteLLM_EndUserTable | None, - update_end_user_table_data: dict[str, object], - prisma_client: "PrismaClient", -) -> None: - """ - Handle object permission updates for customer endpoints. +async def _prepare_customer_object_permission_update( + data: UpdateCustomerRequest, + existing_object_permission_id: str | None, + prisma_client: ObjectPermissionClient, +) -> ObjectPermissionUpsert | None: + if data.object_permission is None: + return None + return await prepare_object_permission_upsert( + new_object_permission=_STR_OBJECT_DICT.validate_python(data.object_permission.model_dump()), + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) - Updates the update_end_user_table_data dict in place with the new object_permission_id. - Args: - non_default_values: Dictionary containing the update values including object_permission - end_user_table_data_typed: Existing end user table data - update_end_user_table_data: Dictionary to update with new object_permission_id - prisma_client: Prisma database client - """ - if "object_permission" in non_default_values: - existing_object_permission_id: Final = ( - end_user_table_data_typed.object_permission_id if end_user_table_data_typed is not None else None +async def _write_customer_budget_update( + tx: ObjectPermissionClient, + budget_data: Mapping[str, object], + existing_budget: LiteLLM_BudgetTable | None, + updated_by: str, +) -> str | None: + if not budget_data: + return None + if existing_budget is not None: + await _typed_table(BudgetRepository(tx)).update( + where={"budget_id": existing_budget.budget_id}, + data=budget_data, ) - object_permission_id: Final = await handle_update_object_permission_common( - data_json=non_default_values, - existing_object_permission_id=existing_object_permission_id, - prisma_client=prisma_client, - ) - if object_permission_id is not None: - update_end_user_table_data["object_permission_id"] = object_permission_id + return None + budget: Final = await _typed_table(BudgetRepository(tx)).create( + data={**budget_data, "created_by": updated_by, "updated_by": updated_by}, + include={"end_users": True}, + ) + return budget.budget_id @router.post( @@ -603,7 +612,6 @@ async def _update_single_customer( try: data_json: Final = _STR_OBJECT_DICT.validate_python(data.json()) - # get the row from db if prisma_client is None: raise Exception("Not connected to DB!") @@ -612,101 +620,81 @@ async def _update_single_customer( k: v for k, v in data_json.items() if _should_update_field(k, v, sent_fields) } - ## Get end user table data ## - end_user_table_data: Final = await _typed_table(EndUserRepository(prisma_client)).find_first( - where={"user_id": data.user_id}, include={"litellm_budget_table": True} - ) - - if end_user_table_data is None: - raise ProxyException( - message=f"End User Id={data.user_id} does not exist in db", - type="not_found", - code=404, - param="user_id", + budget_table_data: Final = { + field: value + for field, value in non_default_values.items() + if field != "budget_id" and field in LiteLLM_BudgetTable.model_fields + } + customer_table_data: Final = { + field: value + for field, value in non_default_values.items() + if field == "budget_id" + or ( + field in LiteLLM_EndUserTable.model_fields + and field not in LiteLLM_BudgetTable.model_fields + and field != "object_permission" ) - - end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) - - ## Get budget table data ## - end_user_budget_table: Final = end_user_table_data_typed.litellm_budget_table - - ## Get all params for budget table ## - budget_table_data: Final = dict[str, object]() - update_end_user_table_data: Final = dict[str, object]() - for k, v in non_default_values.items(): - # budget_id is for linking to existing budget, not for creating new budget - if k == "budget_id": - update_end_user_table_data[k] = v - elif k in LiteLLM_BudgetTable.model_fields: - budget_table_data[k] = v - - elif k in LiteLLM_EndUserTable.model_fields: - update_end_user_table_data[k] = v - - ## Handle object permission updates (MCP servers, vector stores, etc.) - await _handle_customer_object_permission_update( - non_default_values=non_default_values, - end_user_table_data_typed=end_user_table_data_typed, - update_end_user_table_data=update_end_user_table_data, - prisma_client=prisma_client, - ) - - ## Check if we need to create a new budget (only if budget fields are provided, not just budget_id) ## - if budget_table_data: - if end_user_budget_table is None: - ## Create new budget ## - budget_table_data_record = await _typed_table(BudgetRepository(prisma_client)).create( - data={ - **budget_table_data, - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - }, - include={"end_users": True}, + } + async with prisma_client.tx() as transaction: + tx: Final = ObjectPermissionClient(db=transaction) + customer_locked: Final = bool(await transaction.query_raw(_CUSTOMER_UPDATE_LOCK_SQL, data.user_id)) + end_user_table_data: Final = ( + await _typed_table(EndUserRepository(tx)).find_first( + where={"user_id": data.user_id}, include={"litellm_budget_table": True} ) - - update_end_user_table_data["budget_id"] = budget_table_data_record.budget_id - else: - ## Update existing budget ## - budget_table_data_record = await _typed_table(BudgetRepository(prisma_client)).update( - where={"budget_id": end_user_budget_table.budget_id}, - data=budget_table_data, - ) - - ## Update user table, with update params + new budget id (if set) ## - verbose_proxy_logger.debug("/customer/update: Received data = %s", data) - - # Ensure object_permission is not in the update data - # It should have been converted to object_permission_id by handle_update_object_permission_common - if "object_permission" in update_end_user_table_data: - verbose_proxy_logger.warning( - "object_permission still in update_end_user_table_data: %s", - update_end_user_table_data.get("object_permission"), + if customer_locked + else None ) - update_end_user_table_data.pop("object_permission", None) - - if data.user_id is not None and len(data.user_id) > 0: - update_end_user_table_data["user_id"] = data.user_id - verbose_proxy_logger.debug("In update customer, user_id condition block.") - response: Final = await _typed_table(EndUserRepository(prisma_client)).update( + if end_user_table_data is None: + raise ProxyException( + message=f"End User Id={data.user_id} does not exist in db", + type="not_found", + code=404, + param="user_id", + ) + if not data.user_id: + raise ValueError(f"user_id is required, passed user_id = {data.user_id}") + end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) + permission_upsert: Final = await _prepare_customer_object_permission_update( + data=data, + existing_object_permission_id=end_user_table_data_typed.object_permission_id, + prisma_client=tx, + ) + permission: Final = ( + await ObjectPermissionRepository(tx).table.upsert( + where={"object_permission_id": permission_upsert.object_permission_id}, + data={"create": permission_upsert.record, "update": permission_upsert.update_record}, + ) + if permission_upsert is not None + else None + ) + budget_id: Final = await _write_customer_budget_update( + tx=tx, + budget_data=budget_table_data, + existing_budget=end_user_table_data_typed.litellm_budget_table, + updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, + ) + update_data: Final = { + **customer_table_data, + "user_id": data.user_id, + **({"object_permission_id": permission.object_permission_id} if permission is not None else {}), + **({"budget_id": budget_id} if budget_id is not None else {}), + } + response: Final = await _typed_table(EndUserRepository(tx)).update( where={"user_id": data.user_id}, - data=update_end_user_table_data, + data=update_data, include={"litellm_budget_table": True, "object_permission": True}, ) if response is None: raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}") - verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) - - await _evict_end_user_cache_keys(_end_user_cache_keys((data.user_id,))) - - return _to_customer_response(response) - else: - raise ValueError(f"user_id is required, passed user_id = {data.user_id}") - - # update based on remaining passed in values + customer_response: Final = _to_customer_response(response) + await _evict_end_user_cache_keys(_end_user_cache_keys((data.user_id,))) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.update_end_user(): Exception occured - %s", e) return handle_exception_on_proxy(e) + else: + return customer_response @router.post( diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index 2a1445d86d1..fefe549692d 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1,7 +1,8 @@ -from collections.abc import Mapping -from contextlib import contextmanager -from types import MappingProxyType -from typing import Final +import asyncio +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager, contextmanager +from types import MappingProxyType, SimpleNamespace +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -9,9 +10,12 @@ from fastapi import FastAPI, HTTPException, Request, status from fastapi.responses import JSONResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient +from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow from pydantic import TypeAdapter from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_ObjectPermissionBase, LiteLLM_EndUserTable, LiteLLM_ObjectPermissionTable, LiteLLM_UserTable, @@ -25,7 +29,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( end_user_cache_key, end_user_restricted_registry_cache_key, ) -from litellm.proxy.management_endpoints.customer_endpoints import _should_update_field, router +from litellm.proxy.management_endpoints.customer_endpoints import _should_update_field, _update_single_customer, router from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -99,6 +103,20 @@ def test_should_update_field( @pytest.fixture def mock_prisma_client(): with patch("litellm.proxy.proxy_server.prisma_client") as mock: + + async def lock_customer(_query: str, user_id: str) -> tuple[Mapping[str, str], ...]: + customer: Final = mock.db.litellm_endusertable.find_first.return_value + if customer is None: + return () + row: Final = LiteLLM_EndUserTable.model_validate(customer.model_dump()) + return ({"user_id": user_id},) if row.user_id == user_id else () + + mock.configure_mock( + **{ + "db.query_raw": AsyncMock(side_effect=lock_customer), + "tx.return_value.__aenter__": AsyncMock(side_effect=lambda: mock.db), + } + ) yield mock @@ -1203,6 +1221,8 @@ class _CustomerBatchDatabase: self, customers: tuple[LiteLLM_EndUserTable, ...], permissions: tuple[LiteLLM_ObjectPermissionTable, ...] = (), + budgets: tuple[LiteLLM_BudgetTable, ...] = (), + failed_customer_ids: frozenset[str] = frozenset(), ) -> None: self.customers: Mapping[str, LiteLLM_EndUserTable] = MappingProxyType( {customer.user_id: customer for customer in customers} @@ -1210,6 +1230,13 @@ class _CustomerBatchDatabase: self.permissions: Mapping[str, LiteLLM_ObjectPermissionTable] = MappingProxyType( {permission.object_permission_id: permission for permission in permissions} ) + self.budgets: Mapping[str | None, LiteLLM_BudgetTable] = MappingProxyType( + {budget.budget_id: budget for budget in budgets} + ) + self.failed_customer_ids = failed_customer_ids + + async def query_raw(self, _query: str, user_id: str) -> tuple[Mapping[str, str], ...]: + return ({"user_id": user_id},) if user_id in self.customers else () async def find_customer( self, *, where: Mapping[str, str], include: Mapping[str, bool] @@ -1219,6 +1246,10 @@ class _CustomerBatchDatabase: async def update_customer( self, *, where: Mapping[str, str], data: Mapping[str, object], include: Mapping[str, bool] ) -> LiteLLM_EndUserTable: + if where["user_id"] in self.failed_customer_ids: + raise ValueError("Customer write rejected") + if data.get("budget_id") and data["budget_id"] not in self.budgets: + raise ValueError("Budget does not exist") existing: Final = self.customers[where["user_id"]] updated: Final = LiteLLM_EndUserTable.model_validate({**existing.model_dump(), **data}) permission: Final = self.permissions.get(updated.object_permission_id or "") @@ -1232,7 +1263,8 @@ class _CustomerBatchDatabase: async def upsert_permission( self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] ) -> LiteLLM_ObjectPermissionTable: - stored: Final = data["update"] if where["object_permission_id"] in self.permissions else data["create"] + existing: Final = self.permissions.get(where["object_permission_id"]) + stored: Final = {**existing.model_dump(), **data["update"]} if existing is not None else data["create"] tool_permissions: Final = stored.get("mcp_tool_permissions") decoded: Final = ( TypeAdapter(dict[str, list[str]]).validate_json(tool_permissions) @@ -1243,17 +1275,176 @@ class _CustomerBatchDatabase: self.permissions = MappingProxyType({**self.permissions, permission.object_permission_id: permission}) return permission + async def create_budget(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> LiteLLM_BudgetTable: + budget: Final = LiteLLM_BudgetTable.model_validate({**data, "budget_id": "created-budget"}) + self.budgets = MappingProxyType({**self.budgets, budget.budget_id: budget}) + return budget + + async def update_budget(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_BudgetTable: + existing: Final = self.budgets[where["budget_id"]] + budget: Final = LiteLLM_BudgetTable.model_validate({**existing.model_dump(), **data}) + self.budgets = MappingProxyType({**self.budgets, budget.budget_id: budget}) + return budget + def attach(self, prisma_client: MagicMock) -> None: + @asynccontextmanager + async def transaction() -> AsyncIterator[MagicMock]: + customers: Final = self.customers + permissions: Final = self.permissions + budgets: Final = self.budgets + try: + yield prisma_client.db + except Exception: + self.customers = customers + self.permissions = permissions + self.budgets = budgets + raise + prisma_client.configure_mock( **{ + "tx": MagicMock(side_effect=transaction), + "db.query_raw": AsyncMock(side_effect=self.query_raw), "db.litellm_endusertable.find_first": AsyncMock(side_effect=self.find_customer), "db.litellm_endusertable.update": AsyncMock(side_effect=self.update_customer), "db.litellm_objectpermissiontable.find_unique": AsyncMock(side_effect=self.find_permission), "db.litellm_objectpermissiontable.upsert": AsyncMock(side_effect=self.upsert_permission), + "db.litellm_budgettable.create": AsyncMock(side_effect=self.create_budget), + "db.litellm_budgettable.update": AsyncMock(side_effect=self.update_budget), } ) +def test_customer_update_does_not_write_row_created_after_empty_lock_result( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None +) -> None: + appearing: Final = LiteLLM_EndUserTable(user_id="created-after-lock", alias="Original", blocked=False) + database: Final = _CustomerBatchDatabase(()) + database.attach(mock_prisma_client) + + async def lock_missing_customer(_query: str, _user_id: str) -> tuple[Mapping[str, object], ...]: + database.customers = MappingProxyType({appearing.user_id: appearing}) + return () + + @asynccontextmanager + async def transaction() -> AsyncIterator[MagicMock]: + yield mock_prisma_client.db + + mock_prisma_client.configure_mock( + **{ + "tx": MagicMock(side_effect=transaction), + "db.query_raw": AsyncMock(side_effect=lock_missing_customer), + } + ) + with _end_user_cache_doubles() as (recording_cache, mock_publish): + response: Final = client.post( + "/customer/update", + json={"user_id": appearing.user_id, "alias": "Updated"}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 404, response.text + assert response.json()["error"]["type"] == "not_found" + assert database.customers == {appearing.user_id: appearing} + mock_prisma_client.db.litellm_endusertable.find_first.assert_not_awaited() + mock_prisma_client.db.litellm_endusertable.update.assert_not_awaited() + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_awaited() + mock_prisma_client.db.litellm_budgettable.create.assert_not_awaited() + mock_prisma_client.db.litellm_budgettable.update.assert_not_awaited() + assert recording_cache.deleted == [] + assert _published_keys(mock_publish) == [] + + +class _ConcurrentCustomerTransaction: + def __init__( + self, database: _CustomerBatchDatabase, row_lock: asyncio.Lock, unlocked_reads: asyncio.Barrier + ) -> None: + self.database = database + self.row_lock = row_lock + self.unlocked_reads = unlocked_reads + self.has_lock = False + self.litellm_endusertable = SimpleNamespace( + find_first=AsyncMock(side_effect=self.find_customer), + update=AsyncMock(side_effect=database.update_customer), + ) + self.litellm_objectpermissiontable = SimpleNamespace( + find_unique=AsyncMock(side_effect=database.find_permission), + upsert=AsyncMock(side_effect=database.upsert_permission), + ) + self.litellm_mcpservertable = SimpleNamespace(find_many=AsyncMock(return_value=[])) + self.litellm_budgettable = SimpleNamespace( + create=AsyncMock(side_effect=database.create_budget), + update=AsyncMock(side_effect=database.update_budget), + ) + + async def query_raw(self, query: str, user_id: object) -> tuple[Mapping[str, object], ...]: + if query.rstrip().upper().endswith("FOR UPDATE"): + await self.row_lock.acquire() + self.has_lock = True + return ({"user_id": user_id},) if user_id in self.database.customers else () + + async def find_customer( + self, *, where: Mapping[str, str], include: Mapping[str, bool] + ) -> LiteLLM_EndUserTable | None: + snapshot: Final = await self.database.find_customer(where=where, include=include) + if not self.has_lock: + await self.unlocked_reads.wait() + return snapshot + + +@pytest.mark.asyncio +@pytest.mark.parametrize("toolsets_first", (False, True)) +async def test_concurrent_customer_updates_share_new_permission_without_losing_unsent_grants( + mock_prisma_client: MagicMock, toolsets_first: bool +) -> None: + from litellm.proxy._types import UpdateCustomerRequest + + original: Final = LiteLLM_EndUserTable(user_id="customer-id", blocked=True, models=["existing-model"]) + database: Final = _CustomerBatchDatabase((original,)) + row_lock: Final = asyncio.Lock() + unlocked_reads: Final = asyncio.Barrier(2) + + @asynccontextmanager + async def transaction() -> AsyncIterator[_ConcurrentCustomerTransaction]: + tx: Final = _ConcurrentCustomerTransaction(database, row_lock, unlocked_reads) + try: + yield tx + finally: + if tx.has_lock: + row_lock.release() + + mock_prisma_client.configure_mock( + **{ + "db": _ConcurrentCustomerTransaction(database, row_lock, unlocked_reads), + "tx": MagicMock(side_effect=transaction), + } + ) + server_patch: Final = UpdateCustomerRequest( + user_id=original.user_id, object_permission=LiteLLM_ObjectPermissionBase(mcp_servers=["new-server"]) + ) + toolset_patch: Final = UpdateCustomerRequest( + user_id=original.user_id, object_permission=LiteLLM_ObjectPermissionBase(mcp_toolsets=["new-toolset"]) + ) + patches: Final = (toolset_patch, server_patch) if toolsets_first else (server_patch, toolset_patch) + auth: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + results: Final = await asyncio.gather( + *(_update_single_customer(data=customer, user_api_key_dict=auth) for customer in patches) + ) + + assert all(isinstance(result, CustomerResponse) for result in results) + permission_id: Final = database.customers[original.user_id].object_permission_id + assert permission_id is not None + expected_permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id=permission_id, mcp_servers=["new-server"], mcp_toolsets=["new-toolset"] + ) + assert database.permissions == {permission_id: expected_permission} + assert database.customers == { + original.user_id: original.model_copy( + update={"object_permission_id": permission_id, "object_permission": expected_permission} + ) + } + + @pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) def test_customer_bulk_update_applies_each_customers_fields( mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str @@ -1365,6 +1556,158 @@ def test_customer_bulk_update_reports_unknown_middle_item_and_updates_later_cust assert "bulk-missing" not in database.customers +@pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) +@pytest.mark.parametrize("budget_edit", ("link", "create", "update")) +def test_customer_bulk_update_rolls_back_failed_item_permissions_and_budgets( + mock_prisma_client: MagicMock, + mock_user_api_key_auth: None, + route: str, + budget_edit: Literal["link", "create", "update"], +) -> None: + permissions: Final = tuple( + LiteLLM_ObjectPermissionTable( + object_permission_id=f"permission-{suffix}", + mcp_servers=[f"server-{suffix}"], + mcp_toolsets=[f"old-toolset-{suffix}"], + ) + for suffix in ("a", "b", "c") + ) + original_budget: Final = LiteLLM_BudgetTable(budget_id="budget-b", max_budget=10) + budgets: Final = (original_budget,) if budget_edit == "update" else () + customers: Final = tuple( + LiteLLM_EndUserTable( + user_id=f"bulk-{suffix}", + blocked=False, + alias=f"Original {suffix}", + object_permission_id=f"permission-{suffix}", + budget_id=original_budget.budget_id if suffix == "b" and budgets else None, + litellm_budget_table=original_budget if suffix == "b" and budgets else None, + ) + for suffix in ("a", "b", "c") + ) + database: Final = _CustomerBatchDatabase( + customers, + permissions, + budgets, + failed_customer_ids=frozenset() if budget_edit == "link" else frozenset({"bulk-b"}), + ) + database.attach(mock_prisma_client) + failed_budget_patch: Final = {"budget_id": "missing-budget"} if budget_edit == "link" else {"max_budget": 25} + cache: Final = UserApiKeyCache() + for customer in customers: + cache.set_cache(key=end_user_cache_key(customer.user_id), value=f"cached-{customer.user_id}") + + with patch.multiple("litellm.proxy.proxy_server", user_api_key_cache=cache, redis_usage_cache=None): + response: Final = client.post( + route, + json={ + "customers": [ + {"user_id": "bulk-a", "object_permission": {"mcp_toolsets": ["new-toolset"]}}, + { + "user_id": "bulk-b", + "object_permission": {"mcp_toolsets": ["new-toolset"]}, + **failed_budget_patch, + }, + {"user_id": "bulk-c", "object_permission": {"mcp_toolsets": ["new-toolset"]}}, + ] + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200, response.text + body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) + assert (body["total_requested"], body["successful_updates"], body["failed_updates"]) == (3, 2, 1) + results: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_python(body["results"]) + assert tuple((result["user_id"], result["success"]) for result in results) == ( + ("bulk-a", True), + ("bulk-b", False), + ("bulk-c", True), + ) + assert results[1]["error"] == ("Budget does not exist" if budget_edit == "link" else "Customer write rejected") + assert database.customers["bulk-b"] == customers[1] + assert database.permissions == { + "permission-a": permissions[0].model_copy(update={"mcp_toolsets": ["new-toolset"]}), + "permission-b": permissions[1], + "permission-c": permissions[2].model_copy(update={"mcp_toolsets": ["new-toolset"]}), + } + assert database.budgets == {budget.budget_id: budget for budget in budgets} + assert cache.get_cache(key=end_user_cache_key("bulk-a")) is None + assert cache.get_cache(key=end_user_cache_key("bulk-b")) == "cached-bulk-b" + assert cache.get_cache(key=end_user_cache_key("bulk-c")) is None + + +@pytest.mark.parametrize("route", ("/customer/update", "/end_user/update")) +def test_customer_update_rolls_back_created_permission_and_budget_after_customer_write_failure( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + original: Final = LiteLLM_EndUserTable(user_id="customer-id", blocked=True, models=["old-model"]) + database: Final = _CustomerBatchDatabase((original,), failed_customer_ids=frozenset({original.user_id})) + database.attach(mock_prisma_client) + + response: Final = client.post( + route, + json={ + "user_id": original.user_id, + "max_budget": 25, + "object_permission": {"mcp_toolsets": ["new-toolset"]}, + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 500, response.text + assert response.json()["error"]["message"] == "Customer write rejected" + assert database.customers == {original.user_id: original} + assert database.permissions == {} + assert database.budgets == {} + + +@pytest.mark.parametrize("route", ("/customer/update", "/end_user/update")) +def test_customer_update_rolls_back_writes_when_saved_response_cannot_be_validated( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="permission-id", mcp_toolsets=["old-toolset"] + ) + original: Final = LiteLLM_EndUserTable( + user_id="customer-id", blocked=True, models=["old-model"], object_permission_id=permission.object_permission_id + ) + database: Final = _CustomerBatchDatabase((original,), (permission,)) + database.attach(mock_prisma_client) + + async def save_invalid_response( + *, where: Mapping[str, str], data: Mapping[str, object], include: Mapping[str, bool] + ) -> PrismaEndUserRow: + saved: Final = await database.update_customer(where=where, data=data, include=include) + saved_permission: Final = database.permissions[permission.object_permission_id] + return PrismaEndUserRow.model_validate( + { + **saved.model_dump(), + "object_permission": { + **saved_permission.model_dump(), + "mcp_tool_permissions": '{"server-id": 7}', + }, + } + ) + + mock_prisma_client.configure_mock( + **{"db.litellm_endusertable.update": AsyncMock(side_effect=save_invalid_response)} + ) + + response: Final = client.post( + route, + json={ + "user_id": original.user_id, + "models": ["new-model"], + "object_permission": {"mcp_toolsets": ["new-toolset"]}, + }, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 500, response.text + assert database.customers == {original.user_id: original} + assert database.permissions == {permission.object_permission_id: permission} + + @pytest.mark.parametrize("route", ("/customer/bulk_update", "/end_user/bulk_update")) def test_customer_bulk_update_toolsets_only_preserves_each_customers_other_permissions( mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str From d55ab1905dcbaec89880bdcb7ffee75d0fe2d61d Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 15:00:13 -0400 Subject: [PATCH 14/56] refactor(proxy): remove redundant customer update ID guard --- litellm/proxy/management_endpoints/customer_endpoints.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index f47a46ef38d..a624da0aee0 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -652,8 +652,6 @@ async def _update_single_customer( code=404, param="user_id", ) - if not data.user_id: - raise ValueError(f"user_id is required, passed user_id = {data.user_id}") end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) permission_upsert: Final = await _prepare_customer_object_permission_update( data=data, From 211b0c37251fe16d79e2186e74de86ebdb45731c Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 15:04:15 -0400 Subject: [PATCH 15/56] fix(ui): allow editing customer-owned budgets --- .../customers/_components/CustomerForm.tsx | 35 ++++++++++----- .../_components/CustomerInfoView.tsx | 3 ++ .../CustomersPanel.integration.test.tsx | 44 ++++++++++++++++++- .../customers/_components/CustomersPanel.tsx | 6 +++ .../_components/customerPayload.test.ts | 20 ++++++++- .../customers/_components/customerPayload.ts | 6 +++ 6 files changed, 100 insertions(+), 14 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx index ab85dcac7ee..4d959b9c77b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx @@ -57,10 +57,23 @@ const regionOptions = [ { value: "us", label: "US" }, ]; +const getBudgetFieldCopy = (hasCurrentBudget: boolean, isEdit: boolean) => { + if (hasCurrentBudget) { + return { emptyOption: "Keep current budget", description: "Select another budget to change the link" }; + } + return { + emptyOption: "No budget", + description: isEdit + ? "Leaving both budget fields empty keeps the current budget" + : "Choose a budget or enter a maximum below", + }; +}; + type CustomerFormProps = { customer?: EndUser | null; isEdit?: boolean; isSaving?: boolean; + hasSharedBudget?: boolean; onSubmit: (values: CustomerFormValues) => void | Promise; onCancel: () => void; }; @@ -69,29 +82,27 @@ export function CustomerForm({ customer = null, isEdit = false, isSaving = false, + hasSharedBudget, onSubmit, onCancel, }: CustomerFormProps) { const { accessToken } = useAuthorized(); const { data: budgets = [] } = useBudgetOptions(accessToken); const { data: models = [] } = useUserModels(); - const initialValues = toCustomerFormValues(customer); + const initialValues = { ...toCustomerFormValues(customer), ...(isEdit ? { budget_id: "" } : {}) }; const form = useZodForm(isEdit ? editCustomerSchema : customerSchema, { defaultValues: initialValues, }); - const hasLinkedBudget = isEdit && Boolean(initialValues.budget_id); - const budgetFieldCopy = hasLinkedBudget - ? { emptyOption: "Keep current budget", description: "Select another budget to change the link" } - : { - emptyOption: "No budget", - description: isEdit - ? "Leaving both budget fields empty keeps the current budget" - : "Choose a budget or enter a maximum below", - }; + const hasCurrentBudget = isEdit && Boolean(customer?.budget_id); + const budgetFieldCopy = getBudgetFieldCopy(hasCurrentBudget, isEdit); const handleSubmit = form.handleSubmit(async (values) => { - if (hasLinkedBudget && !values.budget_id && values.max_budget.trim()) { + const isMaximumEdit = !values.budget_id && Boolean(values.max_budget.trim()); + if (hasCurrentBudget && hasSharedBudget !== false && isMaximumEdit) { form.setError("max_budget", { - message: "This customer has a linked budget. Leave Max Budget empty or choose another budget", + message: + hasSharedBudget === undefined + ? "Unable to verify budget sharing. Reload customers before changing Max Budget" + : "This customer has a shared budget. Leave Max Budget empty or choose another budget", }); return; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx index 141b89ea9cd..ed4e21bf225 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx @@ -35,6 +35,7 @@ type CustomerInfoViewProps = { isEdit: boolean; tab: CustomerTab; isBlocking: boolean; + hasSharedBudget: boolean | undefined; onTabChange: (tab: CustomerTab) => void; onBack: () => void; onEdit: () => void; @@ -170,6 +171,7 @@ export default function CustomerInfoView({ customer, accessToken, canModify, + hasSharedBudget, isEdit, tab, onTabChange, @@ -255,6 +257,7 @@ export default function CustomerInfoView({ customer={customer} isEdit isSaving={isSaving} + hasSharedBudget={hasSharedBudget} onSubmit={handleSave} onCancel={onCancelEdit} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx index 78aa3511fe5..df7a9a14a35 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx @@ -565,6 +565,48 @@ describe("Customers management", () => { expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); }); + it("updates a customer's own maximum without replacing its budget link", async () => { + const user = userEvent.setup(); + const budget = { budget_id: "own-budget", max_budget: 100, created_at: "2026-01-01T00:00:00Z" }; + const customer: EndUser = { ...activeCustomer, budget_id: budget.budget_id, litellm_budget_table: budget }; + get.mockResolvedValue([budget]); + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [customer, blockedCustomer] : customer, + isLoading: false, + isError: false, + })); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + expect(screen.getByLabelText("Budget")).toHaveTextContent("Keep current budget"); + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "25" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: customer.user_id, alias: "Alice", models: [], max_budget: 25 }, + }), + ); + }); + + it.each(["loading", "error"])("keeps maximum edits guarded while the customer list is %s", async (status) => { + const user = userEvent.setup(); + const budget = { budget_id: "budget-id", max_budget: 100, created_at: "2026-01-01T00:00:00Z" }; + const customer: EndUser = { ...activeCustomer, budget_id: budget.budget_id, litellm_budget_table: budget }; + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? undefined : customer, + isLoading: path === "/customer/list" && status === "loading", + isError: path === "/customer/list" && status === "error", + })); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "25" } }); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect( + await screen.findByText("Unable to verify budget sharing. Reload customers before changing Max Budget"), + ).toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + }); + it("rejects a max budget when keeping a shared budget link", async () => { const user = userEvent.setup(); const budget = { budget_id: "shared-budget", max_budget: 100, created_at: "2026-01-01T00:00:00Z" }; @@ -586,7 +628,7 @@ describe("Customers management", () => { await user.click(screen.getByRole("button", { name: "Save Changes" })); expect( - await screen.findByText("This customer has a linked budget. Leave Max Budget empty or choose another budget"), + await screen.findByText("This customer has a shared budget. Leave Max Budget empty or choose another budget"), ).toBeInTheDocument(); expect(post).not.toHaveBeenCalled(); fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "" } }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx index 74f72c98584..bd9345087d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx @@ -23,6 +23,7 @@ import { all_admin_roles, isProxyAdminRole } from "@/utils/roles"; import { CreateCustomerModal } from "./CreateCustomerModal"; import CustomerInfoView from "./CustomerInfoView"; import CustomersTable from "./CustomersTable"; +import { hasSharedCustomerBudget } from "./customerPayload"; const CUSTOMER_VIEW_PARSERS = { customer: parseAsString, @@ -136,6 +137,11 @@ export default function CustomersPanel() { canModify={canModify} premiumUser={premiumUser} isEdit={edit} + hasSharedBudget={ + customer.data && customers.data && !customers.isError + ? hasSharedCustomerBudget(customer.data, customers.data) + : undefined + } tab={tab} onTabChange={(value) => void setView({ customer_tab: value, edit: null })} isBlocking={isBlocking} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts index d5f8b85c30a..32effab2c38 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "vitest"; import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; -import { buildCustomerPayload, toCustomerFormValues } from "./customerPayload"; +import { buildCustomerPayload, hasSharedCustomerBudget, toCustomerFormValues } from "./customerPayload"; const customer: EndUser = { user_id: "customer-id", blocked: false, spend: 0 }; const permission = { @@ -27,6 +27,24 @@ const server: MCPServer = { updated_by: "admin", }; +describe("shared customer budgets", () => { + it("allows a budget linked to only one customer", () => { + const owned = { ...customer, budget_id: "own-budget" }; + expect( + hasSharedCustomerBudget(owned, [owned, { ...customer, user_id: "another-customer", budget_id: "other-budget" }]), + ).toBe(false); + }); + + it("guards a budget referenced by another customer", () => { + const shared = { ...customer, budget_id: "shared-budget" }; + expect(hasSharedCustomerBudget(shared, [shared, { ...shared, user_id: "another-customer" }])).toBe(true); + }); + + it("does not treat customers without budgets as sharing a budget", () => { + expect(hasSharedCustomerBudget(customer, [customer, { ...customer, user_id: "another-customer" }])).toBe(false); + }); +}); + describe("customer form values", () => { it("starts with empty optional values and grants", () => { const expected = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts index 68ba8fbd75c..1b7c014fafb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts @@ -4,6 +4,12 @@ import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { mcpServersForIdentifier } from "@/components/mcp_server_management/effectiveMcpServers"; import { extractMcpEntitlement, type McpEntitlementUpdate } from "@/components/mcp_server_management/mcpEntitlement"; +export const hasSharedCustomerBudget = (customer: EndUser, customers: readonly EndUser[]): boolean => + Boolean( + customer.budget_id && + customers.some((other) => other.user_id !== customer.user_id && other.budget_id === customer.budget_id), + ); + export type CustomerFormValues = { user_id: string; alias: string; From cdc094f4a28b56c8ccd9cf1787e0348b00819ef5 Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 18:56:18 -0400 Subject: [PATCH 16/56] refactor(proxy): use the shared permission client protocol --- .../object_permission_utils.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index b061f264b47..b6c00af9903 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -21,10 +21,10 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, object_permission_cache_key from litellm.proxy.utils import PrismaClient from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.prisma_protocols import DatabaseClient from litellm.repositories.table_repositories import MCPServerRepository if TYPE_CHECKING: - from prisma import Prisma from prisma import models as prisma_models from litellm.proxy._types import ( @@ -33,11 +33,6 @@ if TYPE_CHECKING: ) -@dataclass(frozen=True, slots=True) -class ObjectPermissionClient: - db: "Prisma" - - async def attach_object_permission_to_dict( data_dict: dict[str, object], prisma_client: PrismaClient, @@ -89,7 +84,7 @@ class ObjectPermissionUpsert: async def prepare_object_permission_upsert( new_object_permission: Mapping[str, object], existing_object_permission_id: str | None, - prisma_client: PrismaClient | ObjectPermissionClient, + prisma_client: DatabaseClient, ) -> ObjectPermissionUpsert: """ Read-and-merge half of an object permission upsert; performs no writes. @@ -126,7 +121,6 @@ async def prepare_object_permission_upsert( for field, value in new_object_permission.items() if value is not None and field != "object_permission_id" } - # Omitted request fields become None; preserve stored grants and Prisma column defaults. merged: Final[dict[str, object]] = { **existing_fields, **requested_fields, @@ -284,7 +278,7 @@ def _mcp_server_identifier_matches(server: object, identifier: str) -> bool: async def _get_db_mcp_servers_by_identifiers( identifiers: AbstractSet[str], - prisma_client: PrismaClient | ObjectPermissionClient | None, + prisma_client: DatabaseClient | None, ) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]": if prisma_client is None or not identifiers: return [] @@ -303,7 +297,7 @@ async def _get_db_mcp_servers_by_identifiers( async def _resolve_mcp_server_identifiers_to_ids( identifiers: AbstractSet[str], - prisma_client: PrismaClient | ObjectPermissionClient | None, + prisma_client: DatabaseClient | None, ) -> dict[str, set[str]]: """ Resolve MCP permission entries written as server_id, alias, or server_name @@ -360,7 +354,7 @@ def _mcp_tool_permission_entries(raw: object) -> Mapping[str, frozenset[str]]: async def reject_ambiguous_mcp_tool_permission_keys( new_mcp_tool_permissions: object, existing_mcp_tool_permissions: object, - prisma_client: PrismaClient | ObjectPermissionClient | None, + prisma_client: DatabaseClient | None, ) -> None: """ A name or alias shared by several MCP servers cannot key ``mcp_tool_permissions``: From 51dfceddf92200e75bfca48913b43fc22ad0e410 Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 18:59:57 -0400 Subject: [PATCH 17/56] fix(proxy): validate customer block-state changes before saving --- litellm/proxy/_types.py | 2 +- .../customer_endpoints.py | 71 +++++++----- .../test_customer_endpoints.py | 103 ++++++++++++++++-- 3 files changed, 138 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 477edf47c5f..fb35c24da38 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1259,7 +1259,7 @@ class ProviderInfo(LiteLLMPydanticObjectBase): class BlockUsers(LiteLLMPydanticObjectBase): - user_ids: list[str] # required + user_ids: list[str] = Field(min_length=1, max_length=1000) class ModelParams(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index a86c792ff6e..dacde45239f 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -24,6 +24,7 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_BudgetTable as PrismaBudgetRow from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList from litellm.proxy.utils import PrismaClient import litellm @@ -45,7 +46,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository -from litellm.repositories.chunked_in import find_many_in +from litellm.repositories.chunked_in import find_many_in, update_many_in from litellm.repositories.table_repositories import EndUserRepository from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, @@ -156,6 +157,29 @@ def _to_customer_response(record: BaseModel) -> CustomerResponse: return CustomerResponse.model_validate(record.model_dump()) +def _configured_blocked_user_hook() -> "ENTERPRISE_BlockedUserList": + 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 + }, + ) + configured_hook: Final = next( + (callback for callback in litellm.callbacks if isinstance(callback, ENTERPRISE_BlockedUserList)), None + ) + if configured_hook is None: + raise HTTPException( + status_code=400, + detail={"error": "Blocked user check was never set. This call has no effect."}, + ) + + return configured_hook + + @router.post( "/end_user/block", tags=["Customer Management"], @@ -187,6 +211,7 @@ async def block_user(data: BlockUsers) -> BlockUsersResponse: """ from litellm.proxy.proxy_server import prisma_client + _configured_blocked_user_hook() try: if prisma_client is None: raise HTTPException( @@ -240,43 +265,33 @@ async def unblock_user(data: BlockUsers) -> UnblockUsersResponse: """ from litellm.proxy.proxy_server import prisma_client - try: - from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList - except ImportError: + configured_hook: Final = _configured_blocked_user_hook() + + requested_ids: Final = frozenset(data.user_ids) + file_ids: Final = ( + _STRING_IDS.validate_python(configured_hook.blocked_user_list or ()) + if isinstance(litellm.blocked_user_list, str) + else () + ) + if requested_ids.intersection(file_ids): 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": "These customers are blocked by a file. Update the file and restart the proxy to unblock them." }, ) - configured_hook: Final = next( - (callback for callback in litellm.callbacks if isinstance(callback, ENTERPRISE_BlockedUserList)), None - ) - if configured_hook is None: - raise HTTPException( - status_code=400, - detail={"error": "Blocked user check was never set. This call has no effect."}, - ) - - requested_ids: Final = frozenset(data.user_ids) - configured_ids: Final = _STRING_IDS.validate_python( - configured_hook.blocked_user_list - if isinstance(litellm.blocked_user_list, str) - else litellm.blocked_user_list or () - ) - remaining_legacy_ids: Final = tuple(user_id for user_id in configured_ids if user_id not in requested_ids) table: Final = _typed_table(EndUserRepository(prisma_client)) if prisma_client is not None else None if table is not None: - await table.update_many(where={"user_id": {"in": data.user_ids}}, data={"blocked": False}) + await update_many_in(table, "user_id", data.user_ids, data={"blocked": False}, atomicity="per_chunk_ok") + configured_ids: Final = ( + file_ids + if isinstance(litellm.blocked_user_list, str) + else _STRING_IDS.validate_python(litellm.blocked_user_list or ()) + ) + remaining_legacy_ids: Final = tuple(user_id for user_id in configured_ids if user_id not in requested_ids) if isinstance(litellm.blocked_user_list, list): litellm.blocked_user_list[:] = remaining_legacy_ids await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) - if isinstance(litellm.blocked_user_list, str) and requested_ids.intersection(configured_ids): - raise HTTPException( - status_code=500, - detail={"error": "`blocked_user_list` must be set as a list. Filepaths can't be updated."}, - ) return UnblockUsersResponse(blocked_users=list(remaining_legacy_ids)) diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index ee8d404c8f5..a4b33081c39 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1,3 +1,5 @@ +import asyncio +from collections.abc import Iterator, Mapping from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -11,6 +13,7 @@ from fastapi.routing import APIRoute from fastapi.testclient import TestClient from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow +import litellm from litellm.proxy._types import ( BlockUsers, LiteLLM_EndUserTable, @@ -101,6 +104,16 @@ def mock_prisma_client(): yield mock +@pytest.fixture +def configured_block_hook(mock_prisma_client: MagicMock) -> Iterator[None]: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + with patch("litellm.blocked_user_list", []): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + with patch("litellm.callbacks", [hook]): + yield + + @pytest.fixture def mock_user_api_key_auth(): original_overrides = app.dependency_overrides.copy() @@ -558,7 +571,7 @@ def test_update_customer_response_keeps_nested_budget_server_fields(mock_prisma_ def test_block_customer_success_serializes_through_response_model( - mock_prisma_client: MagicMock, mock_user_api_key_auth: object + mock_prisma_client: MagicMock, mock_user_api_key_auth: object, configured_block_hook: None ) -> None: """ /customer/block returns {"blocked_users": []}. With @@ -1107,7 +1120,9 @@ def test_customer_update_invalidates_end_user_and_registry_caches(mock_prisma_cl assert _published_keys(mock_publish) == ["end_user_id:c1", "end_user_restricted_registry", "litellm:end_user_id:c1"] -def test_customer_block_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): +def test_customer_block_invalidates_end_user_and_registry_caches( + mock_prisma_client, mock_user_api_key_auth, configured_block_hook: None +): """Blocking is the one mutation that must take effect instantly; a stale registry keeps serving it.""" mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock( return_value=LiteLLM_EndUserTable(user_id="c1", blocked=True) @@ -1138,7 +1153,7 @@ def test_customer_block_invalidates_end_user_and_registry_caches(mock_prisma_cli def test_customer_block_legacy_route_matches_canonical_response( - mock_prisma_client: MagicMock, mock_user_api_key_auth: object + mock_prisma_client: MagicMock, mock_user_api_key_auth: object, configured_block_hook: None ) -> None: blocked_row: Final = PrismaEndUserRow( user_id="legacy-customer", blocked=True, alias="Legacy customer", spend=0, models=[] @@ -1265,7 +1280,8 @@ async def test_customer_block_unblock_refreshes_enterprise_hook_cache(mock_prism cached_row: Final = cache.get_cache(end_user_block_cache_key(active_row.user_id), model_type=PrismaEndUserRow) assert cached_row is not None and cached_row.blocked is False - await block_user(customer_ids) + with patch("litellm.callbacks", [hook]): + await block_user(customer_ids) with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): await hook.async_pre_call_hook(auth, cache, {"user": active_row.user_id}, "completion") @@ -1434,7 +1450,7 @@ async def test_customer_unblock_database_only_hook_refreshes_cached_block(mock_p @pytest.mark.asyncio @pytest.mark.parametrize("in_file", (False, True)) -async def test_customer_unblock_file_list_clears_database_and_keeps_file( +async def test_customer_unblock_file_list_validates_before_database_write( mock_prisma_client: MagicMock, tmp_path: Path, in_file: bool ) -> None: from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList @@ -1460,24 +1476,93 @@ async def test_customer_unblock_file_list_clears_database_and_keeps_file( hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) with patch("litellm.callbacks", [hook]): if in_file: - with pytest.raises(HTTPException, match="Filepaths can't be updated") as error: + with pytest.raises(HTTPException, match="Update the file and restart the proxy") as error: await unblock_user(BlockUsers(user_ids=[user_id])) - assert error.value.status_code == 500 + assert error.value.status_code == 400 else: response: Final = await unblock_user(BlockUsers(user_ids=[user_id])) assert response.blocked_users == ["legacy-only"] - assert cache.get_cache(end_user_block_cache_key(user_id), model_type=PrismaEndUserRow) is None + cached_row: Final = cache.get_cache(end_user_block_cache_key(user_id), model_type=PrismaEndUserRow) if in_file: + assert cached_row is not None and cached_row.blocked is True with pytest.raises(HTTPException, match="User blocked from making LLM API Calls"): await hook.async_pre_call_hook(auth, cache, {"user": user_id}, "completion") else: + assert cached_row is None await hook.async_pre_call_hook(auth, cache, {"user": user_id}, "completion") - update_many.assert_awaited_once_with(where={"user_id": {"in": [user_id]}}, data={"blocked": False}) + if in_file: + update_many.assert_not_awaited() + else: + update_many.assert_awaited_once_with(where={"user_id": {"in": [user_id]}}, data={"blocked": False}) assert block_file.read_text() == file_contents +@pytest.mark.parametrize("route", ("/customer/block", "/end_user/block")) +def test_customer_block_requires_loaded_hook_before_writing( + mock_prisma_client: MagicMock, mock_user_api_key_auth: object, route: str +) -> None: + upsert: Final = AsyncMock() + mock_prisma_client.db = SimpleNamespace(litellm_endusertable=SimpleNamespace(upsert=upsert)) + with patch("litellm.callbacks", []), _end_user_cache_doubles() as (cache, publish): + response: Final = client.post(route, json={"user_ids": ["c1"]}, headers={"Authorization": "Bearer k"}) + assert response.status_code == 400 + assert response.json() == {"detail": {"error": "Blocked user check was never set. This call has no effect."}} + upsert.assert_not_awaited() + assert cache.deleted == [] + publish.assert_not_awaited() + + +@pytest.mark.parametrize("route", ("/customer/block", "/end_user/block", "/customer/unblock", "/end_user/unblock")) +@pytest.mark.parametrize("size", (0, 1001)) +def test_customer_block_requests_reject_empty_or_oversized_batches( + mock_prisma_client: MagicMock, mock_user_api_key_auth: object, route: str, size: int +) -> None: + upsert: Final = AsyncMock() + update_many: Final = AsyncMock() + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace(upsert=upsert, update_many=update_many) + ) + response: Final = client.post( + route, json={"user_ids": [f"c{index}" for index in range(size)]}, headers={"Authorization": "Bearer k"} + ) + assert response.status_code == 422 + upsert.assert_not_awaited() + update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_concurrent_customer_unblocks_keep_both_ids_unblocked(mock_prisma_client: MagicMock) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + first_write_started: Final = asyncio.Event() + release_first_write: Final = asyncio.Event() + + async def write_block_state(where: Mapping[str, object], data: Mapping[str, object]) -> int: + assert data == {"blocked": False} + if where == {"user_id": {"in": ["c1"]}}: + first_write_started.set() + await release_first_write.wait() + return 1 + + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace(update_many=AsyncMock(side_effect=write_block_state)) + ) + with patch("litellm.blocked_user_list", ["c1", "c2", "still-blocked"]), _end_user_cache_doubles(): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=mock_prisma_client) + with patch("litellm.callbacks", [hook]): + first_request: Final = asyncio.create_task(unblock_user(BlockUsers(user_ids=["c1"]))) + await first_write_started.wait() + second_response: Final = await unblock_user(BlockUsers(user_ids=["c2"])) + assert second_response.blocked_users == ["c1", "still-blocked"] + release_first_write.set() + first_response: Final = await first_request + assert first_response.blocked_users == ["still-blocked"] + assert hook.blocked_user_list == ["still-blocked"] + assert litellm.blocked_user_list == ["still-blocked"] + + def test_customer_delete_invalidates_end_user_and_registry_caches(mock_prisma_client, mock_user_api_key_auth): """Without this a deleted customer keeps its cached budget and block enforced until the TTL expires.""" mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock( From 2f8252216bd1d7304a28142dd47b65c8d47e2eec Mon Sep 17 00:00:00 2001 From: lewis Date: Wed, 14 Oct 2026 19:00:11 -0400 Subject: [PATCH 18/56] fix(ui): refine customer controls and MCP permission edits --- .../_components/CustomerBlockConfirmation.tsx | 61 ++++++++++ .../customers/_components/CustomerForm.tsx | 6 +- .../_components/CustomerInfoView.tsx | 9 +- .../CustomersPanel.integration.test.tsx | 112 ++++++++++++++---- .../customers/_components/CustomersPanel.tsx | 28 +++-- .../CustomersTable.integration.test.tsx | 16 ++- .../customers/_components/CustomersTable.tsx | 52 +++++--- .../_components/CustomersTableColumns.test.ts | 56 --------- .../CustomersTableColumns.test.tsx | 54 ++++++++- .../_components/CustomersTableColumns.tsx | 99 ++++++++++------ .../_components/customerCsvExport.test.ts | 13 +- .../_components/customerCsvExport.ts | 2 +- .../_components/customerPayload.test.ts | 23 +++- .../customers/_components/customerPayload.ts | 13 +- .../_components/customerResources.ts | 4 +- .../view_users/user_info_view.test.tsx | 38 +++--- .../mcp_server_management/mcpEntitlement.ts | 60 +++++----- .../components/object_permissions_view.tsx | 3 + .../permissions/MCPServerPermissions.test.tsx | 9 +- .../permissions/MCPServerPermissions.tsx | 8 +- .../shared/table_cells/models_cell.test.tsx | 5 + .../shared/table_cells/models_cell.tsx | 21 ++-- 22 files changed, 476 insertions(+), 216 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBlockConfirmation.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBlockConfirmation.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBlockConfirmation.tsx new file mode 100644 index 00000000000..46cb3e41a37 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerBlockConfirmation.tsx @@ -0,0 +1,61 @@ +"use client"; + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; + +type CustomerBlockConfirmationProps = { + open: boolean; + count: number; + blocked: boolean; + isPending: boolean; + onCancel: () => void; + onConfirm: () => void; +}; + +export function CustomerBlockConfirmation({ + open, + count, + blocked, + isPending, + onCancel, + onConfirm, +}: CustomerBlockConfirmationProps) { + const action = blocked ? "Block" : "Unblock"; + const customerLabel = count === 1 ? "this customer" : "these customers"; + return ( + { + if (!nextOpen && !isPending) onCancel(); + }} + > + + + + {action} {count} customer{count === 1 ? "" : "s"}? + + + {blocked + ? `Requests for ${customerLabel} will be rejected until unblocked. Existing settings and permissions are kept.` + : `Requests for ${customerLabel} will be allowed again, subject to existing limits and permissions.`} + + + + + + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx index 4d959b9c77b..b73479df3a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx @@ -163,7 +163,11 @@ export function CustomerForm({ control={form.control} name="max_budget" label="Max Budget (USD)" - description={isEdit ? "Leave empty to keep the maximum. Updates require more than zero" : undefined} + description={ + isEdit + ? "Leave empty to keep the maximum. Budgets shared with other resources cannot be changed here. Updates require more than zero" + : undefined + } > {({ ref, ...control }) => ( - + {customer.allowed_model_region?.toUpperCase() ?? "Any"} @@ -107,7 +107,11 @@ function CustomerOverview({ customer, accessToken }: Pick
- +
); } @@ -131,6 +135,7 @@ function CustomerSettings({ customer, accessToken }: Pick , action: string, count: number) => { + const dialog = await screen.findByRole("dialog", { name: `${action} ${count} customer${count === 1 ? "" : "s"}?` }); + expect(post).not.toHaveBeenCalled(); + await user.click(within(dialog).getByRole("button", { name: action })); +}; + describe("Customers management", () => { beforeEach(() => { vi.clearAllMocks(); @@ -105,6 +111,7 @@ describe("Customers management", () => { await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); expect(screen.getByText("2 selected")).toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Block selected" })); + await confirmBlocking(user, "Block", 2); await waitFor(() => expect(post).toHaveBeenCalledExactlyOnceWith("/customer/block", { accessToken: "sk-test", @@ -119,13 +126,19 @@ describe("Customers management", () => { const user = userEvent.setup(); renderWithProviders(); await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); - expect(screen.getByRole("button", { name: "Block selected" })).toBeDisabled(); - expect(screen.getByRole("button", { name: "Unblock selected" })).toBeDisabled(); - await user.hover(screen.getByLabelText(`Bulk ${action} requires Enterprise`)); + const button = screen.getByRole("button", { name: `${action} selected` }); + expect(button).toHaveAttribute("aria-disabled", "true"); + await user.tab({ shift: true }); + await user.tab({ shift: true }); + if (action === "Block") await user.tab({ shift: true }); + expect(button).toHaveFocus(); expect(await screen.findByRole("tooltip")).toHaveTextContent( "This is a LiteLLM Enterprise feature, and requires a valid key to use.", ); - await user.click(screen.getByLabelText(`Bulk ${action} requires Enterprise`)); + expect(button).toHaveAccessibleDescription( + "This is a LiteLLM Enterprise feature, and requires a valid key to use.", + ); + await user.keyboard("{Enter}"); expect(post).not.toHaveBeenCalled(); }); @@ -135,6 +148,7 @@ describe("Customers management", () => { renderWithProviders(); await user.click(screen.getByRole("checkbox", { name: `Select customer ${blockedCustomer.user_id}` })); await user.click(screen.getByRole("button", { name: "Unblock selected" })); + await confirmBlocking(user, "Unblock", 1); await waitFor(() => expect(post).toHaveBeenCalledExactlyOnceWith("/customer/unblock", { accessToken: "sk-test", @@ -156,7 +170,9 @@ describe("Customers management", () => { ); await user.click(screen.getByRole("checkbox", { name: `Select customer ${blockedCustomer.user_id}` })); await user.click(screen.getByRole("button", { name: "Unblock selected" })); + await confirmBlocking(user, "Unblock", 1); expect(await screen.findByText(message)).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Cancel" })); expect(screen.getByRole("checkbox", { name: `Select customer ${blockedCustomer.user_id}` })).toBeChecked(); expect(within(screen.getByRole("row", { name: /bob@example.test/ })).getByText("Blocked")).toBeInTheDocument(); expect(screen.queryByText("1 customer unblocked")).not.toBeInTheDocument(); @@ -174,9 +190,7 @@ describe("Customers management", () => { }); renderWithProviders(); expect(within(screen.getByRole("row", { name: /alice@example.test/ })).getByText("< $0.01")).toBeInTheDocument(); - expect( - within(screen.getByRole("row", { name: /bob@example.test/ })).queryByText(/\$0\.00/), - ).not.toBeInTheDocument(); + expect(within(screen.getByRole("row", { name: /bob@example.test/ })).getByText("$0.00")).toBeInTheDocument(); }); it.each([ @@ -204,7 +218,7 @@ describe("Customers management", () => { renderWithProviders(); const row = screen.getByRole("row", { name: /alice@example.test/ }); - expect(within(row).getByText(budgetLabel)).toBeInTheDocument(); + expect(within(row).getByText(maxBudget === null ? "· Unlimited" : `of ${budgetLabel}`)).toBeInTheDocument(); await user.click(within(row).getByRole("button", { name: /Alice.*alice@example.test/ })); expect(await screen.findByText("Budget Status")).toBeInTheDocument(); expect(screen.getAllByText("$0.00")).toHaveLength(maxBudget === 0 ? 2 : 1); @@ -215,16 +229,19 @@ describe("Customers management", () => { }, ); - it("keeps MCP access optional while showing each customer status", async () => { + it("keeps model and MCP details optional while showing each customer status", async () => { const user = userEvent.setup(); renderWithProviders(); expect(screen.queryByRole("columnheader", { name: "MCP Access" })).not.toBeInTheDocument(); + expect(screen.queryByRole("columnheader", { name: "Models" })).not.toBeInTheDocument(); await user.click(screen.getByRole("button", { name: "Columns" })); await user.click(await screen.findByRole("menuitemcheckbox", { name: "MCP Access" })); + await user.click(await screen.findByRole("menuitemcheckbox", { name: "Models" })); await user.keyboard("{Escape}"); const alice = screen.getByRole("row", { name: /alice@example.test/ }); expect(within(alice).getByText("Active")).toBeInTheDocument(); + expect(within(alice).getByText("All models the key can use")).toBeInTheDocument(); expect(within(alice).getByText("1 group, 1 toolset")).toBeInTheDocument(); const bob = screen.getByRole("row", { name: /bob@example.test/ }); expect(within(bob).getByText("Blocked")).toBeInTheDocument(); @@ -247,12 +264,18 @@ describe("Customers management", () => { name: surface === "detail" ? "More customer actions" : `Actions for ${customer.user_id}`, }), ); - expect(await screen.findByRole("menuitem", { name: action })).toHaveAttribute("aria-disabled", "true"); - await user.hover(screen.getByLabelText(`${action} requires Enterprise`)); + const menuAction = await screen.findByRole("menuitem", { name: action }); + expect(menuAction).toHaveAttribute("aria-disabled", "true"); + await user.keyboard("{Home}"); + if (surface === "row") await user.keyboard("{ArrowDown}"); + expect(menuAction).toHaveFocus(); expect(await screen.findByRole("tooltip")).toHaveTextContent( "This is a LiteLLM Enterprise feature, and requires a valid key to use.", ); - await user.click(screen.getByLabelText(`${action} requires Enterprise`)); + expect(menuAction).toHaveAccessibleDescription( + "This is a LiteLLM Enterprise feature, and requires a valid key to use.", + ); + await user.keyboard("{Enter}"); expect(post).not.toHaveBeenCalled(); }); @@ -277,6 +300,7 @@ describe("Customers management", () => { const menuAction = await screen.findByRole("menuitem", { name: action }); expect(menuAction).not.toHaveAttribute("aria-disabled", "true"); await user.click(menuAction); + await confirmBlocking(user, action, 1); await waitFor(() => expect(post).toHaveBeenCalledExactlyOnceWith(endpoint, { accessToken: "sk-test", @@ -285,8 +309,37 @@ describe("Customers management", () => { ); }); + it.each(["row", "detail", "bulk"])("allows cancelling a block from the %s surface", async (surface) => { + authorization.premiumUser = true; + const user = userEvent.setup(); + renderWithProviders( + , + surface === "detail" ? { searchParams: "customer=alice%40example.test" } : undefined, + ); + if (surface === "bulk") { + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + await user.click(screen.getByRole("button", { name: "Block selected" })); + } else { + await user.click( + screen.getByRole("button", { + name: surface === "detail" ? "More customer actions" : `Actions for ${activeCustomer.user_id}`, + }), + ); + await user.click(await screen.findByRole("menuitem", { name: "Block" })); + } + const count = surface === "bulk" ? 2 : 1; + const dialog = await screen.findByRole("dialog", { name: `Block ${count} customer${count === 1 ? "" : "s"}?` }); + await user.click(within(dialog).getByRole("button", { name: "Cancel" })); + expect(post).not.toHaveBeenCalled(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + it.each([null, activeCustomer.user_id])("shows the customer ID once when alias is %s", (alias) => { - useQuery.mockReturnValue({ data: { ...activeCustomer, alias }, isLoading: false, isError: false }); + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [{ ...activeCustomer, alias }] : { ...activeCustomer, alias }, + isLoading: false, + isError: false, + })); renderWithProviders(, { searchParams: "customer=alice%40example.test" }); expect(screen.getAllByText(activeCustomer.user_id)).toHaveLength(1); expect(screen.getByRole("heading", { name: /alice@example.test/ })).toBeInTheDocument(); @@ -369,7 +422,11 @@ describe("Customers management", () => { mcp_toolsets: [], }, }; - useQuery.mockReturnValue({ data: directCustomer, isLoading: false, isError: false }); + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [directCustomer] : directCustomer, + isLoading: false, + isError: false, + })); fetchServers.mockResolvedValue([ { server_id: "server-1", @@ -419,7 +476,11 @@ describe("Customers management", () => { vector_stores: [], }, }; - useQuery.mockReturnValue({ data: directCustomer, isLoading: false, isError: false }); + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [directCustomer] : directCustomer, + isLoading: false, + isError: false, + })); fetchServers.mockRejectedValue(new Error("Catalog unavailable")); renderWithProviders( <> @@ -494,7 +555,7 @@ describe("Customers management", () => { expect(screen.queryByRole("button", { name: "Save Changes" })).not.toBeInTheDocument(); expect(post).not.toHaveBeenCalled(); }); - it("matches the Organizations money, model and rate-limit columns", () => { + it("uses the shared spend/budget cell and hides model details by default", () => { const customer: EndUser = { ...activeCustomer, budget_id: "shared-budget", @@ -509,16 +570,16 @@ describe("Customers management", () => { useQuery.mockReturnValue({ data: [customer], isLoading: false, isError: false }); renderWithProviders(); - expect(screen.getByText("Manage end-user spend, budgets, rate limits, model and MCP access")).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: "Spend (USD)" })).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: "Budget (USD)" })).toBeInTheDocument(); + expect(screen.getByText("Manage customer spend, budgets, rate limits, model and MCP access")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Spend / Budget" })).toBeInTheDocument(); + expect(screen.queryByRole("columnheader", { name: "Models" })).not.toBeInTheDocument(); expect(screen.getByRole("columnheader", { name: "TPM / RPM Limits" })).toBeInTheDocument(); const row = screen.getByRole("row", { name: /alice@example.test/ }); expect(within(row).getByText("$12.50")).toBeInTheDocument(); - expect(within(row).getByText("$100.00")).toBeInTheDocument(); + expect(within(row).getByText("of $100.00")).toBeInTheDocument(); expect(within(row).getByText("TPM: 600")).toBeInTheDocument(); expect(within(row).getByText("RPM: 6")).toBeInTheDocument(); - expect(within(row).getByText("All Proxy Models")).toBeInTheDocument(); + expect(within(row).queryByText("All models the key can use")).not.toBeInTheDocument(); }); it("shows Overview cards and shared permissions before entering Settings", async () => { @@ -536,7 +597,11 @@ describe("Customers management", () => { created_at: "2026-01-01T00:00:00Z", }, }; - useQuery.mockReturnValue({ data: customer, isLoading: false, isError: false }); + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [customer] : customer, + isLoading: false, + isError: false, + })); renderWithProviders(, { searchParams: "customer=alice%40example.test", onUrlUpdate }); expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); @@ -545,7 +610,7 @@ describe("Customers management", () => { expect(screen.getByText("Reset: monthly")).toBeInTheDocument(); expect(screen.getByText("Rate Limits")).toBeInTheDocument(); expect(screen.getByText("Object Permissions")).toBeInTheDocument(); - expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); + expect(screen.getByText("All models the key can use")).toBeInTheDocument(); expect(screen.getByText("EU")).toBeInTheDocument(); expect(screen.getByText("default-model")).toBeInTheDocument(); expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); @@ -661,6 +726,7 @@ describe("Customers management", () => { }), ); await user.click(await screen.findByRole("menuitem", { name: "Unblock" })); + await confirmBlocking(user, "Unblock", 1); expect(await screen.findByText(message)).toBeInTheDocument(); expect(screen.getByText("Request Error")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx index bd9345087d3..06757881732 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx @@ -2,7 +2,7 @@ import { Contact, Plus } from "lucide-react"; import { parseAsBoolean, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; -import { useCallback, useState } from "react"; +import { useCallback, useMemo, useState } from "react"; import { useCustomer, @@ -20,10 +20,11 @@ import { Button } from "@/components/ui/button"; import { toast } from "@/lib/toast"; import { all_admin_roles, isProxyAdminRole } from "@/utils/roles"; +import { CustomerBlockConfirmation } from "./CustomerBlockConfirmation"; import { CreateCustomerModal } from "./CreateCustomerModal"; import CustomerInfoView from "./CustomerInfoView"; import CustomersTable from "./CustomersTable"; -import { hasSharedCustomerBudget } from "./customerPayload"; +import { getCustomerBudgetCounts, hasSharedCustomerBudget } from "./customerPayload"; const CUSTOMER_VIEW_PARSERS = { customer: parseAsString, @@ -38,7 +39,7 @@ function CustomersHeader({ canModify, onCreate }: { canModify: boolean; onCreate Customers - Manage end-user spend, budgets, rate limits, model and MCP access + Manage customer spend, budgets, rate limits, model and MCP access {canModify && ( + ); - if (premiumUser) return button; + if (premiumUser) + return ( + + ); return ( - }> - {button} + } + > + {content} - + This is a LiteLLM Enterprise feature, and requires a valid key to use. @@ -110,6 +120,7 @@ export default function CustomersTable({ onColumnFiltersChange, } = useUrlTableState(TABLE_STATE_OPTIONS); const [filtersOpen, setFiltersOpen] = useState(false); + const [blockTarget, setBlockTarget] = useState<{ userIds: string[]; blocked: boolean } | null>(null); const [rowSelection, setRowSelection] = useState({}); const filterScope = JSON.stringify([search, columnFilters]); const [selectionScope, setSelectionScope] = useState(filterScope); @@ -140,11 +151,24 @@ export default function CustomersTable({ const handleBulkSetBlocked = async (userIds: string[], blocked: boolean) => { const canSetBlocked = canModify && !isBlocking && premiumUser; if (!userIds.length || !canSetBlocked) return; - if (await onBulkSetBlocked(userIds, blocked)) setRowSelection({}); + if (await onBulkSetBlocked(userIds, blocked)) { + setRowSelection({}); + setBlockTarget(null); + } }; return (
+ setBlockTarget(null)} + onConfirm={() => { + if (blockTarget) void handleBulkSetBlocked(blockTarget.userIds, blockTarget.blocked); + }} + /> row.original), ); - toast.success(`Exported ${count} customers`); + toast.success(`Exported ${count} customer${count === 1 ? "" : "s"}`); } catch (error) { toast.fromError(error); } @@ -211,13 +235,13 @@ export default function CustomersTable({ blocked premiumUser={premiumUser} disabled={isBlocking} - onClick={() => void handleBulkSetBlocked(selectedIds, true)} + onClick={() => setBlockTarget({ userIds: selectedIds, blocked: true })} /> void handleBulkSetBlocked(selectedIds, false)} + onClick={() => setBlockTarget({ userIds: selectedIds, blocked: false })} />