diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py index f6998d58cd5..c1c57705095 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,14 @@ 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.cache_pydantic_utils import CacheCodec +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 @@ -58,6 +60,13 @@ class ENTERPRISE_BlockedUserList(CustomLogger): if litellm.set_verbose is True: print(print_statement) # noqa + def _database(self) -> Optional[PrismaClient]: + if self.prisma_client is not None: + return self.prisma_client + from litellm.proxy.proxy_server import prisma_client + + return prisma_client + @with_service_target(AUTH_OBJECTS_TARGET) async def async_pre_call_hook( self, @@ -88,14 +97,15 @@ class ENTERPRISE_BlockedUserList(CustomLogger): }, ) - cache_key = f"litellm:end_user_id:{user}" - end_user_cache_obj: Optional[LiteLLM_EndUserTable] = cache.get_cache( # type: ignore - key=cache_key + cache_key: Final = end_user_block_cache_key(f"{user}") + end_user_cache_obj: Final = CacheCodec.deserialize( + cache.get_cache(key=cache_key), model_type=LiteLLM_EndUserTable ) - if end_user_cache_obj is None and self.prisma_client is not None: + database: Final = self._database() + if end_user_cache_obj is None and database is not None: # check db end_user_obj = ( - await self.prisma_client.db.litellm_endusertable.find_unique( + await database.db.litellm_endusertable.find_unique( where={"user_id": user} ) ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000100_customer_budget_reference_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000100_customer_budget_reference_index/migration.sql new file mode 100644 index 00000000000..cadc07c0ec2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000100_customer_budget_reference_index/migration.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_EndUserTable_budget_id_idx" ON "LiteLLM_EndUserTable" ("budget_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000200_key_budget_reference_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000200_key_budget_reference_index/migration.sql new file mode 100644 index 00000000000..88af995aa61 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000200_key_budget_reference_index/migration.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_VerificationToken_budget_id_idx" ON "LiteLLM_VerificationToken" ("budget_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000300_key_customer_budget_reference_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000300_key_customer_budget_reference_index/migration.sql new file mode 100644 index 00000000000..68c99aa75d3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261018000300_key_customer_budget_reference_index/migration.sql @@ -0,0 +1 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_VerificationToken_end_user_budget_id_idx" ON "LiteLLM_VerificationToken" USING GIN (("metadata"->'end_user_budget_id') jsonb_path_ops); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 83a0e6f9e05..cf6ccf5c631 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -563,6 +563,7 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) + @@index([budget_id]) } model LiteLLM_JWTKeyMapping { @@ -678,6 +679,8 @@ model LiteLLM_EndUserTable { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) blocked Boolean @default(false) + + @@index([budget_id]) } // Budget and shared spend for a model access group. The groups themselves are not rows anywhere: 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 4090666e16f..19f615440b8 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 @@ -5,7 +5,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass from datetime import datetime, timezone -from types import MappingProxyType +from types import EllipsisType, MappingProxyType from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import HTTPException @@ -464,6 +464,7 @@ def _admission_failure_fallback( class MCPServerAccess: server_ids: tuple[str, ...] scope: Literal["unscoped", "scoped", "unresolved"] = "unscoped" + end_user_denies_mcp: bool = False @dataclass(frozen=True, slots=True) @@ -1830,7 +1831,21 @@ class MCPRequestHandler: Returns: List[str]: List of allowed MCP servers by server id """ - from litellm.proxy.proxy_server import general_settings + from litellm.proxy.proxy_server import general_settings, prisma_client + + try: + end_user_permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if user_api_key_auth and user_api_key_auth.end_user_id and prisma_client is not None + else None + ) + except UnloadableEntitlementError as error: + verbose_logger.warning("Denying MCP access, customer entitlement unreadable: %s", error) + return MCPServerAccess(server_ids=(), scope="scoped", end_user_denies_mcp=True) + if end_user_permission is not None and SpecialMCPServerNames.no_mcp_servers.value in ( + end_user_permission.mcp_servers or [] + ): + return MCPServerAccess(server_ids=(), scope="scoped", end_user_denies_mcp=True) if managed_agent_policy(user_api_key_auth) is not None: from litellm.proxy._experimental.mcp_server.auth.managed_agent_access import managed_agent_servers @@ -1911,7 +1926,7 @@ class MCPRequestHandler: ######################################################### if user_api_key_auth and user_api_key_auth.end_user_id: allowed_mcp_servers_for_end_user: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( - user_api_key_auth + user_api_key_auth, object_permission=end_user_permission ) # If end_user has explicit MCP server permissions, apply intersection @@ -1925,11 +1940,10 @@ class MCPRequestHandler: # Always apply intersection: key/team AND end_user # This ensures end_user can only access servers that both they AND their key/team are authorized for - filtered_servers: Final = [] - for _mcp_server in allowed_mcp_servers: - if _mcp_server in allowed_mcp_servers_for_end_user: - filtered_servers.append(_mcp_server) - allowed_mcp_servers = filtered_servers + end_user_servers: Final = frozenset(allowed_mcp_servers_for_end_user) + allowed_mcp_servers = [ + server_id for server_id in allowed_mcp_servers if server_id in end_user_servers + ] verbose_logger.debug( "Applied end_user intersection filter. Final allowed servers: %s", allowed_mcp_servers ) @@ -3428,6 +3442,8 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_end_user( user_api_key_auth: UserAPIKeyAuth | None = None, + *, + object_permission: LiteLLM_ObjectPermissionTable | None | EllipsisType = ..., ) -> list[str]: """ Get allowed MCP servers for an end user. @@ -3445,9 +3461,19 @@ 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) - if object_permission is None: + permission: Final = ( + await MCPRequestHandler._get_end_user_object_permission(user_api_key_auth, prisma_client) + if object_permission is ... + else object_permission + ) + if permission is None: return [] + if SpecialMCPServerNames.no_mcp_servers.value in (permission.mcp_servers or []): + return [SpecialMCPServerNames.no_mcp_servers.value] + + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions( + permission, requires_fresh_policy=user_api_key_auth.requires_fresh_policy + ) try: # Permission entries may be server_ids OR names/aliases — expand to ids. @@ -3455,25 +3481,28 @@ 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(permission.mcp_servers or []) # Get MCP servers from access groups access_group_servers: Final = await MCPRequestHandler.get_mcp_servers_from_access_groups( - object_permission.mcp_access_groups or [], + permission.mcp_access_groups or [], requires_fresh_policy=bool(user_api_key_auth and user_api_key_auth.requires_fresh_policy), ) # servers referenced in tool permissions should also be accessible tool_perm_servers: Final = list( - global_mcp_server_manager.expand_tool_permissions(object_permission.mcp_tool_permissions).keys() + global_mcp_server_manager.expand_tool_permissions(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, 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/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 96cf6ee11aa..04b69390505 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3693,6 +3693,14 @@ class MCPServerManager: ) try: + customer_access: Final = ( + access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + if user_api_key_auth and user_api_key_auth.end_user_id + else None + ) + if customer_access is not None and customer_access.end_user_denies_mcp: + return [] + # Admin view with no explicit object permission and no entitlement ceiling resolves the # whole registry, for keys AND admitted session subjects alike (one predicate owns the # question). Seeded into the union rather than returned early so the session resource @@ -3701,7 +3709,7 @@ class MCPServerManager: resolved_access: Final = ( MCPServerAccess(server_ids=()) if admin_unscoped - else access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + else customer_access or access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) ) resolved_server_ids: Final = ( set(self.get_registry().keys()) if admin_unscoped else set(resolved_access.server_ids) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4f77ac0005..3763c8183ff 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1157,7 +1157,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): @@ -1846,6 +1846,20 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): models: list[str] | None = None object_permission: LiteLLM_ObjectPermissionBase | None = None + @model_validator(mode="after") + def validate_budget_update(self) -> "UpdateCustomerRequest": + if {"budget_id", "max_budget"}.issubset(self.model_fields_set): + raise ValueError("Set either 'max_budget' or 'budget_id', not both.") + return self + + +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): """ 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/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index f5e8304c83c..4361f00971b 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", @@ -934,6 +936,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/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index e64fd755ef9..9a1eb26108a 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..8043e787729 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -6,21 +6,24 @@ All /customer management endpoints /customer/new /customer/info /customer/update +/customer/bulk_update /customer/delete """ #### END-USER/CUSTOMER MANAGEMENT #### from collections.abc import Mapping, Sequence from collections.abc import Set as AbstractSet +from dataclasses import dataclass 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 from pydantic import BaseModel, TypeAdapter if TYPE_CHECKING: + from prisma import Prisma from prisma.models import LiteLLM_BudgetTable as PrismaBudgetRow from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow @@ -32,6 +35,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, ) @@ -39,32 +43,58 @@ from litellm.proxy.management_endpoints.common_daily_activity import get_daily_a 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 _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.chunked_in import find_many_in, update_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, ) from litellm.types.proxy.management_endpoints.customer_endpoints import ( BlockUsersResponse, + BulkUpdateCustomerResponse, CustomerResponse, + CustomerUpdateResult, DeleteCustomersResponse, UnblockUsersResponse, ) + +@dataclass(frozen=True, slots=True) +class _TransactionClient: + db: "Prisma" + + _RowT_co: Final = TypeVar("_RowT_co", covariant=True) _STR_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) _CLEARABLE_LIST_FIELDS: Final = frozenset({"models"}) +_CUSTOMER_BUDGET_LOCK_SQL: Final = 'SELECT "budget_id" FROM "LiteLLM_BudgetTable" WHERE "budget_id" = $1 FOR UPDATE' +_CUSTOMER_SHARED_BUDGET_SQL: Final = """ +SELECT "budget_id" FROM "LiteLLM_BudgetTable" WHERE "budget_id" = $1 AND EXISTS ( + SELECT 1 FROM "LiteLLM_EndUserTable" WHERE "budget_id" = $1 AND "user_id" <> $2 + UNION ALL SELECT 1 FROM "LiteLLM_VerificationToken" WHERE "budget_id" = $1 + OR (("metadata"->'end_user_budget_id') @> to_jsonb($1::text) + AND jsonb_typeof("metadata"->'end_user_budget_id') = 'string') + UNION ALL SELECT 1 FROM "LiteLLM_OrganizationTable" WHERE "budget_id" = $1 + UNION ALL SELECT 1 FROM "LiteLLM_ProjectTable" WHERE "budget_id" = $1 + UNION ALL SELECT 1 FROM "LiteLLM_TagTable" WHERE "budget_id" = $1 + UNION ALL SELECT 1 FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE "budget_id" = $1 + UNION ALL SELECT 1 FROM "LiteLLM_TeamMembership" WHERE "budget_id" = $1 + UNION ALL SELECT 1 FROM "LiteLLM_OrganizationMembership" WHERE "budget_id" = $1 + UNION ALL SELECT 1 FROM "LiteLLM_TeamTable" WHERE "metadata"->>'team_member_budget_id' = $1 +) +""" +_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: if value is None: return False - if field in sent_fields and (isinstance(value, bool) or field in _CLEARABLE_LIST_FIELDS): + if field in sent_fields and (isinstance(value, (int, float)) or field in _CLEARABLE_LIST_FIELDS): return True if isinstance(value, (list, dict)) and not value: return False @@ -105,6 +135,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 +167,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: @@ -148,6 +184,34 @@ def _to_customer_response(record: BaseModel) -> CustomerResponse: return CustomerResponse.model_validate(record.model_dump()) +class _BlockedUserHook(Protocol): + @property + def blocked_user_list(self) -> Sequence[str] | None: ... + + +def _configured_blocked_user_hook() -> _BlockedUserHook: + 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"], @@ -160,7 +224,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 @@ -179,26 +243,28 @@ async def block_user(data: BlockUsers): """ from litellm.proxy.proxy_server import prisma_client + _configured_blocked_user_hook() 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 +282,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,38 +295,35 @@ async def unblock_user(data: BlockUsers): }' ``` """ - try: - from enterprise.enterprise_hooks.blocked_user_list import ( - ENTERPRISE_BlockedUserList, - ) - except ImportError: - raise HTTPException( - status_code=400, - detail={ - "error": "Blocked user check was never set. This call has no effect." - + CommonProxyErrors.missing_enterprise_package_docker.value - }, - ) + from litellm.proxy.proxy_server import prisma_client - if ( - not any(isinstance(x, ENTERPRISE_BlockedUserList) for x in litellm.callbacks) - or litellm.blocked_user_list is None + configured_hook: Final = _configured_blocked_user_hook() + + requested_ids: Final = frozenset(data.user_ids) + if isinstance(litellm.blocked_user_list, str) and requested_ids.intersection( + configured_hook.blocked_user_list or () ): raise HTTPException( status_code=400, - detail={"error": "Blocked user check was never set. This call has no effect."}, + detail={ + "error": "These customers are blocked by a file. Update the file and restart the proxy to unblock them." + }, ) - + if prisma_client is not None: + await update_many_in( + _typed_table(EndUserRepository(prisma_client)), + "user_id", + data.user_ids, + data={"blocked": False}, + atomicity="per_chunk_ok", + ) + remaining_legacy_ids: Final = tuple( + user_id for user_id in configured_hook.blocked_user_list or () if user_id not in requested_ids + ) if isinstance(litellm.blocked_user_list, list): - for id in data.user_ids: - litellm.blocked_user_list.remove(id) - else: - raise HTTPException( - status_code=500, - detail={"error": "`blocked_user_list` must be set as a list. Filepaths can't be updated."}, - ) - - return {"blocked_users": litellm.blocked_user_list} + litellm.blocked_user_list[:] = remaining_legacy_ids + await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + return UnblockUsersResponse(blocked_users=list(remaining_legacy_ids)) def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: @@ -289,34 +352,48 @@ 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. - - 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 _customer_budget_update_error( + tx: _TransactionClient, + budget_id: str, + user_id: str, +) -> ProxyException | None: + budget_locked: Final = bool(await tx.db.query_raw(_CUSTOMER_BUDGET_LOCK_SQL, budget_id)) + if not budget_locked: + return ProxyException(message="Budget does not exist", type="not_found", code=404, param="budget_id") + shared: Final = budget_id == litellm.max_end_user_budget_id or bool( + await tx.db.query_raw(_CUSTOMER_SHARED_BUDGET_SQL, budget_id, user_id) + ) + if shared: + return ProxyException( + message=( + "Cannot update max_budget for a shared budget. Update the budget directly or assign a different budget." + ), + type="bad_request", + code=400, + param="max_budget", ) - 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, + return None + + +async def _write_customer_budget_update( + tx: _TransactionClient, + 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, ) - 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( @@ -593,6 +670,108 @@ async def end_user_info( raise handle_exception_on_proxy(e) +async def _update_single_customer( + data: UpdateCustomerRequest, + user_api_key_dict: UserAPIKeyAuth, +) -> CustomerResponse | ProxyException: + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + + try: + data_json: Final = _STR_OBJECT_DICT.validate_python(data.json()) + if prisma_client is None: + raise Exception("Not connected to DB!") + + sent_fields: Final = data.fields_set() + non_default_values: Final[dict[str, object]] = { + k: v for k, v in data_json.items() if _should_update_field(k, v, sent_fields) + } + + 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" + ) + } + async with prisma_client.tx() as transaction: + tx: Final = _TransactionClient(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} + ) + if customer_locked + else None + ) + 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", + ) + end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) + budget_error: Final = ( + await _customer_budget_update_error(tx, end_user_table_data_typed.budget_id, data.user_id) + if budget_table_data and end_user_table_data_typed.budget_id is not None + else None + ) + if budget_error is not None: + return budget_error + permission_upsert: Final = ( + await prepare_object_permission_upsert( + new_object_permission=_STR_OBJECT_DICT.validate_python(data.object_permission.model_dump()), + existing_object_permission_id=end_user_table_data_typed.object_permission_id, + prisma_client=tx, + ) + if data.object_permission is not None + else None + ) + 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_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}") + 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( "/customer/update", tags=["Customer Management"], @@ -610,7 +789,7 @@ async def update_end_user( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> CustomerResponse: """ - Example curl + Example curl Parameters: - user_id: str @@ -629,12 +808,13 @@ async def update_end_user( 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: ``` @@ -663,114 +843,75 @@ async def update_end_user( ``` """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + result: Final = await _update_single_customer(data=data, user_api_key_dict=user_api_key_dict) + if isinstance(result, ProxyException): + raise result + return result - 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!") - sent_fields: Final = data.fields_set() - non_default_values: Final[dict[str, object]] = { - k: v for k, v in data_json.items() if _should_update_field(k, v, sent_fields) - } +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) - ## 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", - ) +@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 with the same fields as `/customer/update`. Only proxy admins can call it. - end_user_table_data_typed: Final = LiteLLM_EndUserTable.model_validate(end_user_table_data.model_dump()) + Each item runs in its own locked transaction, so a failed item rolls back only its own writes and + the rest still apply. Results come back in request order with `user_id`, `success`, and either + `updated_customer` or `error`. - ## Get budget table data ## - end_user_budget_table: Final = end_user_table_data_typed.litellm_budget_table + Omitted or null permission fields keep their stored values, supplied lists replace that field, + and empty lists or objects clear it. - ## 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 + `max_budget` only updates a budget used by that customer alone. Shared and default budgets must be + edited through the budget endpoints. Sending both `budget_id` and `max_budget` in any item, even + as null, fails the whole request with HTTP 422. - 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}, - ) - - 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"), - ) - 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( - where={"user_id": data.user_id}, - data=update_end_user_table_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 - - 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) + 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/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 022eba499a4..b1ef8b2022e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2866,7 +2866,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 d7a44cdd0b8..3c75d78e903 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -989,7 +989,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 3468396c6f6..739c1ee9169 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -22,6 +22,7 @@ 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: @@ -78,19 +79,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: DatabaseClient, ) -> 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 @@ -114,9 +117,14 @@ 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" + } merged: Final[dict[str, object]] = { **existing_fields, - **new_object_permission, + **requested_fields, "object_permission_id": object_permission_id, } record: Final[dict[str, object]] = { @@ -127,7 +135,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, + ) _DECODED_JSON: Final = TypeAdapter(object) @@ -181,7 +197,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, }, ) @@ -269,7 +285,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: DatabaseClient | None, ) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]": if prisma_client is None or not identifiers: return [] @@ -288,7 +304,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: DatabaseClient | None, ) -> dict[str, set[str]]: """ Resolve MCP permission entries written as server_id, alias, or server_name @@ -345,7 +361,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: DatabaseClient | None, ) -> None: """ A name or alias shared by several MCP servers cannot key ``mcp_tool_permissions``: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 83a0e6f9e05..cf6ccf5c631 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -563,6 +563,7 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) + @@index([budget_id]) } model LiteLLM_JWTKeyMapping { @@ -678,6 +679,8 @@ model LiteLLM_EndUserTable { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) blocked Boolean @default(false) + + @@index([budget_id]) } // Budget and shared spend for a model access group. The groups themselves are not rows anywhere: 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/schema.prisma b/schema.prisma index 83a0e6f9e05..cf6ccf5c631 100644 --- a/schema.prisma +++ b/schema.prisma @@ -563,6 +563,7 @@ model LiteLLM_VerificationToken { // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 @@index([budget_reset_at, expires]) + @@index([budget_id]) } model LiteLLM_JWTKeyMapping { @@ -678,6 +679,8 @@ model LiteLLM_EndUserTable { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) blocked Boolean @default(false) + + @@index([budget_id]) } // Budget and shared spend for a model access group. The groups themselves are not rows anywhere: diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 00f8f3b6031..94e815325f8 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -94,6 +94,7 @@ GET /guardrails/submissions GET /guardrails/submissions/{guardrail_id} POST /credentials/migrate-encryption POST /customer/block +POST /customer/bulk_update POST /customer/unblock POST /guardrails/apply_guardrail POST /guardrails/register diff --git a/tests/integration/authorization/test_end_user_block.py b/tests/integration/authorization/test_end_user_block.py index 7b1eb8da74a..43b42087593 100644 --- a/tests/integration/authorization/test_end_user_block.py +++ b/tests/integration/authorization/test_end_user_block.py @@ -139,12 +139,13 @@ def test_block_route_marks_the_customer_blocked_and_update_unblocks(gateway: Gat ) -def _assert_unseen_customer_is_blocked(gateway: Gateway, block_route: str) -> None: - with gateway.scenario() as scenario: +def _assert_unseen_customer_is_blocked(gateway: Gateway, tmp_path: Path, block_route: str) -> None: + with ExitStack() as stack: + proxy: Final = _one_callback_proxy(gateway, tmp_path, stack) end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" - scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + stack.callback(proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]}) - blocked: Final = gateway.request("POST", block_route, {"user_ids": [end_user]}) + blocked: Final = proxy.request("POST", block_route, {"user_ids": [end_user]}) assert blocked.status_code == 200, blocked.text blocked_response: Final = _BlockUsersResponse.model_validate_json(blocked.content) assert blocked_response == _BlockUsersResponse( @@ -153,26 +154,26 @@ def _assert_unseen_customer_is_blocked(gateway: Gateway, block_route: str) -> No assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [ {"blocked": True} ] - info: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user}) + info: Final = proxy.request("GET", "/customer/info", params={"end_user_id": end_user}) assert info.status_code == 200, info.text assert _BlockedCustomerInfo.model_validate(info.json()).blocked is True - updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False}) + updated: Final = proxy.request("POST", "/customer/update", {"user_id": end_user, "blocked": False}) assert updated.status_code == 200, updated.text assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [ {"blocked": False} ] - info_after_update: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user}) + info_after_update: Final = proxy.request("GET", "/customer/info", params={"end_user_id": end_user}) assert info_after_update.status_code == 200, info_after_update.text assert _BlockedCustomerInfo.model_validate(info_after_update.json()).blocked is False -def test_end_user_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None: - _assert_unseen_customer_is_blocked(gateway, "/end_user/block") +def test_end_user_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway, tmp_path: Path) -> None: + _assert_unseen_customer_is_blocked(gateway, tmp_path, "/end_user/block") -def test_customer_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None: - _assert_unseen_customer_is_blocked(gateway, "/customer/block") +def test_customer_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway, tmp_path: Path) -> None: + _assert_unseen_customer_is_blocked(gateway, tmp_path, "/customer/block") def test_listed_end_user_is_refused_and_unblock_restores_serving(gateway: Gateway, tmp_path: Path) -> None: @@ -212,7 +213,7 @@ def test_listed_end_user_is_refused_and_unblock_restores_serving(gateway: Gatewa def test_blocked_end_user_is_refused_on_every_proxy_and_unblock_restores( gateway: Gateway, tmp_path: Path, block_route: str ) -> None: - pytest.skip("BUG: blocked end-user chat requests return HTTP 200 instead of HTTP 400") + pytest.skip("BUG: another proxy keeps a customer's previous block state for up to 60 seconds after unblock") with ( ExitStack() as stack, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, @@ -292,7 +293,7 @@ def test_blocked_end_user_is_refused_on_every_proxy_and_unblock_restores( assert _observed(upstream) == expected_observations -def test_unblock_route_requires_the_blocked_user_check_callback(gateway: Gateway, tmp_path: Path) -> None: +def test_block_and_unblock_require_the_blocked_user_check_callback(gateway: Gateway, tmp_path: Path) -> None: with ExitStack() as stack: callback_proxy: Final = _one_callback_proxy(gateway, tmp_path, stack) end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" @@ -310,21 +311,26 @@ def test_unblock_route_requires_the_blocked_user_check_callback(gateway: Gateway ) assert _blocked_row(end_user) == [{"blocked": True}] + unseen_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + shared_block: Final = gateway.request("POST", "/customer/block", {"user_ids": [unseen_user]}) + assert shared_block.status_code == 400, shared_block.text + assert shared_block.json() == shared.json(), shared_block.text + assert _blocked_row(unseen_user) == [] + callback_response: Final = callback_proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]}) - assert callback_response.status_code == 400, callback_response.text - assert callback_response.json() == { - "detail": {"error": "Blocked user check was never set. This call has no effect."} - }, callback_response.text - assert _blocked_row(end_user) == [{"blocked": True}] + assert callback_response.status_code == 200, callback_response.text + assert callback_response.json() == {"blocked_users": []}, callback_response.text + assert _blocked_row(end_user) == [{"blocked": False}] -def test_block_route_returns_the_blocked_records(gateway: Gateway) -> None: +def test_block_route_returns_the_blocked_records(gateway: Gateway, tmp_path: Path) -> None: end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" - with gateway.scenario() as scenario: - created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user}) + with ExitStack() as stack: + proxy: Final = _one_callback_proxy(gateway, tmp_path, stack) + created: Final = proxy.request("POST", "/customer/new", {"user_id": end_user}) assert created.status_code == 200, created.text - scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) - blocked: Final = gateway.request("POST", "/customer/block", {"user_ids": [end_user]}) + stack.callback(proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + blocked: Final = proxy.request("POST", "/customer/block", {"user_ids": [end_user]}) assert blocked.status_code == 200, blocked.text payload: Final = object_value(blocked.json()) blocked_users: Final = payload["blocked_users"] diff --git a/tests/integration/database/migrations/test_customer_budget_reference_lookup.py b/tests/integration/database/migrations/test_customer_budget_reference_lookup.py new file mode 100644 index 00000000000..16c33bdef22 --- /dev/null +++ b/tests/integration/database/migrations/test_customer_budget_reference_lookup.py @@ -0,0 +1,89 @@ +import os +from pathlib import Path +from typing import Final + +import psycopg +from psycopg import sql +from pydantic import BaseModel, Field, TypeAdapter + +from litellm.proxy.management_endpoints.customer_endpoints import _CUSTOMER_SHARED_BUDGET_SQL + +_REPOSITORY: Final = Path(__file__).resolve().parents[4] +_MIGRATIONS: Final = _REPOSITORY / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" +_INDEX_MIGRATIONS: Final = ( + "20261018000100_customer_budget_reference_index", + "20261018000200_key_budget_reference_index", + "20261018000300_key_customer_budget_reference_index", +) +_REFERENCE_TABLES: Final = ( + "LiteLLM_OrganizationTable", + "LiteLLM_ProjectTable", + "LiteLLM_TagTable", + "LiteLLM_ModelAccessGroupBudgetTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationMembership", + "LiteLLM_TeamTable", +) + + +class _PlanNode(BaseModel, frozen=True): + rows: float = Field(default=0, alias="Actual Rows") + removed: float = Field(default=0, alias="Rows Removed by Filter") + loops: float = Field(default=0, alias="Actual Loops") + plans: tuple["_PlanNode", ...] = Field(default=(), alias="Plans") + + def examined_rows(self) -> float: + return (self.rows + self.removed) * self.loops + sum(plan.examined_rows() for plan in self.plans) + + +class _QueryPlan(BaseModel, frozen=True): + plan: _PlanNode = Field(alias="Plan") + + +def test_customer_budget_guard_keeps_generic_plan_work_bounded() -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute('CREATE TABLE "LiteLLM_BudgetTable" (budget_id text PRIMARY KEY)') + connection.execute('CREATE TABLE "LiteLLM_EndUserTable" (user_id text PRIMARY KEY, budget_id text)') + connection.execute( + 'CREATE TABLE "LiteLLM_VerificationToken" (token text PRIMARY KEY, budget_id text, metadata jsonb)' + ) + for table in _REFERENCE_TABLES: + connection.execute( + sql.SQL("CREATE TABLE {} (budget_id text, metadata jsonb)").format(sql.Identifier(table)) + ) + connection.execute("INSERT INTO \"LiteLLM_BudgetTable\" SELECT 'budget-'||n FROM generate_series(0,500) n") + connection.execute( + "INSERT INTO \"LiteLLM_EndUserTable\" SELECT 'customer-'||n,'budget-'||n FROM generate_series(0,500) n" + ) + key_count: Final = 50_000 + connection.execute( + "INSERT INTO \"LiteLLM_VerificationToken\" SELECT 'key-'||n,NULL," + "jsonb_build_object('end_user_budget_id','unrelated-budget') FROM generate_series(1,%s) n", + (key_count,), + ) + for migration in _INDEX_MIGRATIONS: + connection.execute((_MIGRATIONS / migration / "migration.sql").read_text()) + connection.execute("ANALYZE") + connection.execute("SET plan_cache_mode = force_generic_plan") + connection.execute("PREPARE guard(text,text) AS " + _CUSTOMER_SHARED_BUDGET_SQL) + plan_row: Final = connection.execute( + "EXPLAIN (ANALYZE, FORMAT JSON) EXECUTE guard('budget-0','customer-0')" + ).fetchone() + assert plan_row is not None + plan: Final = TypeAdapter(tuple[_QueryPlan, ...]).validate_python(plan_row[0])[0].plan + assert plan.examined_rows() < key_count / 100 + assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == [] + for metadata in ('{"end_user_budget_id": ["budget-0"]}', '{"end_user_budget_id": null}'): + connection.execute( + "INSERT INTO \"LiteLLM_VerificationToken\" VALUES ('malformed',NULL,%s::jsonb)", (metadata,) + ) + assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == [] + connection.execute("DELETE FROM \"LiteLLM_VerificationToken\" WHERE token='malformed'") + connection.execute( + 'INSERT INTO "LiteLLM_VerificationToken" VALUES ' + "('metadata-budget',NULL,jsonb_build_object('end_user_budget_id','budget-0'))" + ) + assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == [("budget-0",)] + connection.execute("DELETE FROM \"LiteLLM_VerificationToken\" WHERE token='metadata-budget'") + connection.execute("INSERT INTO \"LiteLLM_VerificationToken\" VALUES ('direct-budget','budget-0','{}')") + assert connection.execute("EXECUTE guard('budget-0','customer-0')").fetchall() == [("budget-0",)] 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 c4414003ca5..324f4a4f4e9 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 @@ -13,7 +13,7 @@ from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers -from litellm.caching.caching import DualCache +from litellm.caching.dual_cache import DualCache from litellm.constants import MCP_ALL_TOOLS_WILDCARD from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( @@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( _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, LiteLLM_TeamTable, @@ -32,6 +33,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( @@ -10656,26 +10659,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 @@ -10686,6 +10695,207 @@ 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_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", + 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"]}) diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py index d47b2f73a48..2629b1ced7a 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -20427,3 +20427,114 @@ async def test_legacy_upstream_elicitation_rejects_modern_downstream_without_con finally: active_mcp_request_ctx_var.reset(request_token) legacy_server.active_mcp_session_var.reset(session_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["manager", "mcp", "rest"]) +@pytest.mark.parametrize("admin", [False, True]) +@pytest.mark.parametrize("permission_state", ["denied", "stale-grants", "unreadable"]) +async def test_customer_denial_blocks_operator_open_servers( + monkeypatch: pytest.MonkeyPatch, + surface: Literal["manager", "mcp", "rest"], + admin: bool, + permission_state: Literal["denied", "stale-grants", "unreadable"], +) -> None: + from starlette.requests import Request + + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations, rest_endpoints + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="operator-open", name="operator_open", transport=MCPTransport.http, allow_all_keys=True + ) + manager.config_mcp_servers = {server.server_id: server} + permission: Final = LiteLLM_ObjectPermissionTable( + object_permission_id="customer-denial", + mcp_servers=["no-mcp-servers"], + mcp_tool_permissions={server.server_id: ["delete"]} if permission_state == "stale-grants" else {}, + mcp_toolsets=["stale-toolset"] if permission_state == "stale-grants" else [], + mcp_access_groups=["stale-group"] if permission_state == "stale-grants" else [], + ) + cache: Final = DualCache() + await cache.async_set_cache( + key=end_user_cache_key("denied-customer"), + value=LiteLLM_EndUserTable( + user_id="denied-customer", + blocked=False, + object_permission_id=permission.object_permission_id, + object_permission=None if permission_state == "unreadable" else permission, + ), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id="denied-customer", + user_role=LitellmUserRoles.PROXY_ADMIN if admin else None, + object_permission=None if admin else LiteLLM_ObjectPermissionTable( + object_permission_id="key-grant", mcp_servers=[server.server_id] + ), + ).model_copy(update={"via_virtual_key": True}) + _catalog_database(monkeypatch, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + monkeypatch.setattr(operations, "global_mcp_server_manager", manager) + monkeypatch.setattr(rest_endpoints, "global_mcp_server_manager", manager) + + if surface == "manager": + assert await manager.get_allowed_mcp_servers(auth) == [] + elif surface == "mcp": + assert await operations._get_allowed_mcp_servers(auth, None) == [] + else: + request: Final = Request({"type": "http", "method": "POST", "path": "/mcp-rest/tools/call", "headers": []}) + with pytest.raises(HTTPException) as denied: + await rest_endpoints._resolve_allowed_mcp_servers_with_ip_filter(request, auth, server.server_id) + assert denied.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("customer", ["absent", "unrestricted", "empty-group", "limited"]) +@pytest.mark.parametrize("admin", [False, True]) +async def test_customer_without_explicit_deny_keeps_operator_open_servers( + monkeypatch: pytest.MonkeyPatch, + customer: Literal["absent", "unrestricted", "empty-group", "limited"], + admin: bool, +) -> None: + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key + + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="operator-open", name="operator_open", transport=MCPTransport.http, allow_all_keys=True + ) + manager.config_mcp_servers = {server.server_id: server} + cache: Final = DualCache() + await cache.async_set_cache( + key=end_user_cache_key("unrestricted-customer"), + value=LiteLLM_EndUserTable( + user_id="unrestricted-customer", + blocked=False, + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="ordinary-customer", + mcp_servers=["another-server"] if customer == "limited" else [], + mcp_access_groups=["empty-group"] if customer == "empty-group" else [], + ), + ), + ) + auth: Final = UserAPIKeyAuth( + api_key="test-key", + end_user_id=None if customer == "absent" else "unrestricted-customer", + user_role=LitellmUserRoles.PROXY_ADMIN if admin else None, + ).model_copy(update={"via_virtual_key": True}) + _catalog_database(monkeypatch, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(mcp_server_manager, "global_mcp_server_manager", manager) + + assert await manager.get_allowed_mcp_servers(auth) == [server.server_id] 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_budget.py b/tests/unit/proxy/management_endpoints/test_customer_budget.py index 0beca0c15e8..abe96e30b74 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_budget.py +++ b/tests/unit/proxy/management_endpoints/test_customer_budget.py @@ -47,6 +47,15 @@ def mock_budget_table(): return MagicMock(spec=LiteLLM_BudgetTable) +def _run_in_locked_transaction(prisma_client: MagicMock) -> None: + prisma_client.configure_mock( + **{ + "db.query_raw": AsyncMock(return_value=({"user_id": "test-user"},)), + "tx.return_value.__aenter__": AsyncMock(return_value=prisma_client.db), + } + ) + + @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.prisma_client") @patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -69,6 +78,7 @@ async def test_update_customer_with_budget_id( mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( return_value=mock_existing_customer ) + _run_in_locked_transaction(mock_prisma_client) mock_updated_user = MagicMock() mock_updated_user.model_dump.return_value = { @@ -125,6 +135,7 @@ async def test_update_customer_creates_budget_with_proper_relations( mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( return_value=mock_existing_customer ) + _run_in_locked_transaction(mock_prisma_client) # Mock budget creation mock_created_budget = MagicMock() @@ -183,6 +194,7 @@ async def test_update_customer_creates_budget_with_required_fields( mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( return_value=mock_existing_customer ) + _run_in_locked_transaction(mock_prisma_client) # Mock budget creation mock_created_budget = MagicMock() @@ -248,6 +260,7 @@ async def test_update_customer_budget_creation_with_fallback_admin( mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( return_value=mock_existing_customer ) + _run_in_locked_transaction(mock_prisma_client) # Mock budget creation mock_created_budget = MagicMock() @@ -283,67 +296,6 @@ async def test_update_customer_budget_creation_with_fallback_admin( assert creation_data["updated_by"] == "admin" -@pytest.mark.asyncio -@patch("litellm.proxy.proxy_server.prisma_client") -@patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") -async def test_update_customer_with_budget_id_and_creation_fields( - mock_prisma_client, mock_user_api_key_dict, mock_existing_customer -): - """ - Test customer update when both budget_id and budget creation fields are provided. - - When both linking (budget_id) and creation fields are provided, the system - should prioritize creating a new budget and assign its ID to the customer. - """ - # Arrange - mock_existing_customer.model_dump.return_value = { - "user_id": "test-user", - "blocked": False, - "litellm_budget_table": None, - } - - mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( - return_value=mock_existing_customer - ) - - # Mock budget creation - mock_created_budget = MagicMock() - mock_created_budget.budget_id = "new-budget-combo" - mock_prisma_client.db.litellm_budgettable.create = AsyncMock( - return_value=mock_created_budget - ) - - # Mock end user update - mock_updated_user = MagicMock() - mock_updated_user.model_dump.return_value = {"user_id": "test-user", "blocked": False} - mock_prisma_client.db.litellm_endusertable.update = AsyncMock( - return_value=mock_updated_user - ) - - # Create update request with both budget_id and budget creation fields - update_request = UpdateCustomerRequest( - user_id="test-user", - budget_id="existing-budget-link", # For linking to existing budget - max_budget=300.0, # This should trigger new budget creation - rpm_limit=500, # Use valid budget field - ) - - # Act - await update_end_user(update_request, mock_user_api_key_dict) - - # Assert - # Verify budget creation occurred (because max_budget was provided) - mock_prisma_client.db.litellm_budgettable.create.assert_called_once() - - # Verify end user update was called - mock_prisma_client.db.litellm_endusertable.update.assert_called_once() - call_args = mock_prisma_client.db.litellm_endusertable.update.call_args - - # The update data should contain budget_id from the created budget, not the original budget_id - update_data = call_args[1]["data"] - assert update_data["budget_id"] == "new-budget-combo" # From created budget - - def test_new_budget_request_sets_budget_reset_at_when_duration_provided(): """ Test that new_budget_request auto-populates budget_reset_at when diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index 23cb328aabf..2b2585a5bab 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -1,4 +1,12 @@ -from contextlib import contextmanager +import asyncio +import json +import sqlite3 +import sys +from collections.abc import AsyncIterator, Iterator, Mapping +from contextlib import asynccontextmanager, contextmanager +from pathlib import Path +from types import MappingProxyType, SimpleNamespace +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6,15 +14,36 @@ from fastapi import FastAPI, HTTPException, Request, status from fastapi.responses import JSONResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient -from pydantic import BaseModel +from prisma.models import LiteLLM_EndUserTable as PrismaEndUserRow +from pydantic import BaseModel, TypeAdapter +import litellm from litellm.proxy._types import ( + BlockUsers, + CommonProxyErrors, + LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_ObjectPermissionBase, + 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.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, + _update_single_customer, + block_user, + router, + unblock_user, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -56,7 +85,8 @@ client = TestClient(app) ("object_permission", {"mcp_servers": ["s1"]}, frozenset({"object_permission"}), True), ("metadata", ["m1"], frozenset(), True), ("models", ["m1"], frozenset(), True), - ("max_budget", 0, frozenset({"max_budget"}), False), + ("max_budget", 0, frozenset({"max_budget"}), True), + ("max_budget", 0, frozenset(), False), ("max_budget", 5.0, frozenset(), True), ("alias", "a", frozenset(), True), ), @@ -72,6 +102,7 @@ client = TestClient(app) "nonempty-list-other-field", "nonempty-models", "zero-budget", + "omitted-zero-budget", "nonzero-budget", "alias", ), @@ -88,9 +119,33 @@ 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 +@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() @@ -547,25 +602,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, configured_block_hook: None +) -> 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): @@ -916,9 +973,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)) @@ -1007,6 +1062,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"] == {"mcp_toolsets": ["toolset-id"]} + 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=[ @@ -1074,8 +1165,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): @@ -1093,11 +1184,13 @@ 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): +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) @@ -1115,15 +1208,496 @@ 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_success_serializes_db_model_instance(mock_prisma_client, mock_user_api_key_auth): +def test_customer_block_legacy_route_matches_canonical_response( + 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=[] + ) + 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 + + 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") + + 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( + query_raw=AsyncMock(return_value=[{"user_id": before_row.user_id}]), + 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_blocked_user_hook_built_before_database_connects_rejects_blocked_customer( + mock_prisma_client: MagicMock, +) -> None: + from enterprise.enterprise_hooks.blocked_user_list import ENTERPRISE_BlockedUserList + + blocked_row: Final = PrismaEndUserRow(user_id="config-loaded-customer", blocked=True, spend=0, models=[]) + read_customer: Final = AsyncMock(return_value=blocked_row) + mock_prisma_client.db = SimpleNamespace(litellm_endusertable=SimpleNamespace(find_unique=read_customer)) + with patch("litellm.blocked_user_list", None): + hook: Final = ENTERPRISE_BlockedUserList(prisma_client=None) + cache: Final = UserApiKeyCache() + for _ in range(2): + with pytest.raises(HTTPException, match="User blocked from making LLM API Calls") as error: + await hook.async_pre_call_hook( + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + cache, + {"user": blocked_row.user_id}, + "completion", + ) + assert error.value.status_code == 400 + read_customer.assert_awaited_once_with(where={"user_id": 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_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 + + 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="Update the file and restart the proxy") as error: + await unblock_user(BlockUsers(user_ids=[user_id])) + assert error.value.status_code == 400 + else: + response: Final = await unblock_user(BlockUsers(user_ids=[user_id])) + assert response.blocked_users == ["legacy-only"] + + 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") + + 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", "/customer/unblock")) +def test_customer_block_without_enterprise_package_rejects_before_writing( + mock_prisma_client: MagicMock, mock_user_api_key_auth: object, route: str +) -> None: + upsert: Final = AsyncMock() + update_many: Final = AsyncMock() + mock_prisma_client.db = SimpleNamespace( + litellm_endusertable=SimpleNamespace(upsert=upsert, update_many=update_many) + ) + with ( + patch.dict(sys.modules, {"enterprise.enterprise_hooks.blocked_user_list": None}), + _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." + + CommonProxyErrors.missing_enterprise_package_docker.value + } + } + upsert.assert_not_awaited() + update_many.assert_not_awaited() + assert cache.deleted == [] + publish.assert_not_awaited() + + +@pytest.mark.usefixtures("configured_block_hook") +def test_customer_block_without_database_fails_without_evicting_caches(mock_user_api_key_auth: object) -> None: + with patch("litellm.proxy.proxy_server.prisma_client", None), _end_user_cache_doubles() as (cache, publish): + response: Final = client.post( + "/customer/block", json={"user_ids": ["c1"]}, headers={"Authorization": "Bearer k"} + ) + assert response.status_code == 500 + assert "Postgres DB Not connected" in response.text + 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_block_success_serializes_db_model_instance( + mock_prisma_client, mock_user_api_key_auth, configured_block_hook: None +): class MockPrismaEndUser(BaseModel): user_id: str blocked: bool @@ -1168,9 +1742,957 @@ 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", ] + + +class _CustomerBatchDatabase: + def __init__( + self, + customers: tuple[LiteLLM_EndUserTable, ...], + permissions: tuple[LiteLLM_ObjectPermissionTable, ...] = (), + budgets: tuple[LiteLLM_BudgetTable, ...] = (), + failed_customer_ids: frozenset[str] = frozenset(), + budget_references: tuple[tuple[str, str], ...] = (), + ) -> 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} + ) + self.budgets: Mapping[str | None, LiteLLM_BudgetTable] = MappingProxyType( + {budget.budget_id: budget for budget in budgets} + ) + self.failed_customer_ids = failed_customer_ids + self.budget_references = budget_references + self.locked_budget_ids: frozenset[str] = frozenset() + + async def query_raw(self, query: str, identifier: str, *parameters: str) -> tuple[Mapping[str, object], ...]: + if 'FROM "LiteLLM_EndUserTable"' in query and query.rstrip().endswith("FOR UPDATE"): + return ({"user_id": identifier},) if identifier in self.customers else () + if query.rstrip().endswith("FOR UPDATE"): + self.locked_budget_ids = self.locked_budget_ids | {identifier} + return ({"budget_id": identifier},) if identifier in self.budgets else () + assert identifier in self.locked_budget_ids + return self.budget_reference_rows(query, (identifier, *parameters)) + + def budget_reference_rows(self, query: str, parameters: tuple[str, ...]) -> tuple[Mapping[str, object], ...]: + table_names: Final = ( + "LiteLLM_BudgetTable", + "LiteLLM_EndUserTable", + "LiteLLM_VerificationToken", + "LiteLLM_OrganizationTable", + "LiteLLM_ProjectTable", + "LiteLLM_TagTable", + "LiteLLM_ModelAccessGroupBudgetTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationMembership", + "LiteLLM_TeamTable", + ) + with sqlite3.connect(":memory:") as database: + database.row_factory = sqlite3.Row + for table in table_names: + database.execute(f'CREATE TABLE "{table}" (budget_id TEXT, user_id TEXT, metadata TEXT)') + database.executemany( + 'INSERT INTO "LiteLLM_BudgetTable" (budget_id) VALUES (?)', + tuple((budget_id,) for budget_id in self.budgets), + ) + database.executemany( + 'INSERT INTO "LiteLLM_EndUserTable" (budget_id, user_id) VALUES (?, ?)', + tuple((customer.budget_id, customer.user_id) for customer in self.customers.values()), + ) + for reference, budget_id in self.budget_references: + table: Final = ( + "LiteLLM_VerificationToken" + if reference == "key_default" + else "LiteLLM_TeamTable" + if reference == "team_default" + else reference + ) + assert table in table_names + metadata_key: Final = "end_user_budget_id" if reference == "key_default" else "team_member_budget_id" + metadata: Final = json.dumps({metadata_key: budget_id}) if reference.endswith("_default") else "{}" + database.execute( + f'INSERT INTO "{table}" (budget_id, metadata) VALUES (?, ?)', + (None if reference.endswith("_default") else budget_id, metadata), + ) + return TypeAdapter(tuple[dict[str, object], ...]).validate_python( + [ + dict(row) + for row in database.execute( + query.replace("@> to_jsonb($1::text)", "= json_quote($1)") + .replace("jsonb_typeof", "json_type") + .replace("= 'string'", "= 'text'"), + {str(index): value for index, value in enumerate(parameters, 1)}, + ).fetchall() + ] + ) + + 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: + 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 "") + saved: Final = updated.model_copy(update={"object_permission": permission}) + self.customers = MappingProxyType({**self.customers, saved.user_id: saved}) + budget: Final = self.budgets.get(saved.budget_id) + return ( + CustomerResponse.model_validate( + { + **saved.model_dump(), + "litellm_budget_table": {**budget.model_dump(), "created_at": "2026-01-01T00:00:00Z"}, + } + ) + if budget is not None + else 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: + 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) + 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 + + 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 +) -> 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")) +@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 +) -> 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() + + +_SHARED_BUDGET_REFERENCES: Final = ( + "customer", + "LiteLLM_VerificationToken", + "LiteLLM_OrganizationTable", + "LiteLLM_ProjectTable", + "LiteLLM_TagTable", + "LiteLLM_ModelAccessGroupBudgetTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationMembership", + "key_default", + "team_default", + "proxy_default", +) + + +@pytest.mark.parametrize( + "route", ("/customer/update", "/end_user/update", "/customer/bulk_update", "/end_user/bulk_update") +) +@pytest.mark.parametrize("reference", _SHARED_BUDGET_REFERENCES) +def test_customer_budget_update_rejects_shared_references_without_partial_writes( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str, reference: str +) -> None: + budget: Final = LiteLLM_BudgetTable(budget_id="shared-budget", max_budget=50) + permission: Final = LiteLLM_ObjectPermissionTable(object_permission_id="shared-permission", mcp_toolsets=["old"]) + original: Final = LiteLLM_EndUserTable( + user_id="shared-customer", + blocked=False, + alias="Original", + budget_id=budget.budget_id, + litellm_budget_table=budget, + object_permission_id=permission.object_permission_id, + ) + other: Final = LiteLLM_EndUserTable( + user_id="other-customer", blocked=False, budget_id=budget.budget_id if reference == "customer" else None + ) + references: Final = () if reference in {"customer", "proxy_default"} else ((reference, "shared-budget"),) + database: Final = _CustomerBatchDatabase((original, other), (permission,), (budget,), budget_references=references) + database.attach(mock_prisma_client) + update: Final = { + "user_id": original.user_id, + "max_budget": 100, + "alias": "Changed", + "object_permission": {"mcp_toolsets": ["new"]}, + } + bulk: Final = route.endswith("bulk_update") + payload: Final = {"customers": [update, {"user_id": other.user_id, "alias": "Independent"}]} if bulk else update + with patch("litellm.max_end_user_budget_id", "shared-budget" if reference == "proxy_default" else None): + response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"}) + + expected_error: Final = ( + "Cannot update max_budget for a shared budget. Update the budget directly or assign a different budget." + ) + assert response.status_code == (200 if bulk else 400), response.text + if bulk: + body: Final = response.json() + assert (body["total_requested"], body["successful_updates"], body["failed_updates"]) == (2, 1, 1) + assert body["results"][0] == { + "user_id": original.user_id, + "success": False, + "updated_customer": None, + "error": expected_error, + } + assert body["results"][1]["success"] is True + assert database.customers[other.user_id].alias == "Independent" + else: + assert response.json()["error"]["message"] == expected_error + assert database.customers[other.user_id] == other + assert database.customers[original.user_id] == original + assert database.permissions == {permission.object_permission_id: permission} + assert database.budgets == {budget.budget_id: budget} + assert database.locked_budget_ids == {budget.budget_id} + + +@pytest.mark.parametrize("new_max_budget", (100, 0)) +@pytest.mark.parametrize( + "route", ("/customer/update", "/end_user/update", "/customer/bulk_update", "/end_user/bulk_update") +) +def test_customer_budget_update_allows_own_budget_with_unrelated_shared_references( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str, new_max_budget: int +) -> None: + own_budget: Final = LiteLLM_BudgetTable(budget_id="own-budget", max_budget=50) + other_budget: Final = LiteLLM_BudgetTable(budget_id="other-budget", max_budget=75) + own_customer: Final = LiteLLM_EndUserTable( + user_id="own-customer", + blocked=False, + budget_id=own_budget.budget_id, + litellm_budget_table=own_budget, + ) + other: Final = LiteLLM_EndUserTable(user_id="other-customer", blocked=False, budget_id=other_budget.budget_id) + references: Final = tuple( + (reference, "other-budget") + for reference in _SHARED_BUDGET_REFERENCES + if reference not in {"customer", "proxy_default"} + ) + database: Final = _CustomerBatchDatabase( + (own_customer, other), budgets=(own_budget, other_budget), budget_references=references + ) + database.attach(mock_prisma_client) + update: Final = {"user_id": own_customer.user_id, "max_budget": new_max_budget} + payload: Final = {"customers": [update]} if route.endswith("bulk_update") else update + with patch("litellm.max_end_user_budget_id", "other-budget"): + response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == 200, response.text + assert database.budgets[own_budget.budget_id].max_budget == new_max_budget + assert database.budgets[other_budget.budget_id] == other_budget + assert database.customers[own_customer.user_id].budget_id == own_budget.budget_id + assert database.locked_budget_ids == {own_budget.budget_id} + + +@pytest.mark.parametrize( + "route", ("/customer/update", "/end_user/update", "/customer/bulk_update", "/end_user/bulk_update") +) +def test_customer_budget_update_rejects_missing_linked_budget_without_writing( + mock_prisma_client: MagicMock, mock_user_api_key_auth: None, route: str +) -> None: + original: Final = LiteLLM_EndUserTable(user_id="customer", alias="Original", blocked=False, budget_id="gone") + database: Final = _CustomerBatchDatabase((original,)) + database.attach(mock_prisma_client) + update: Final = {"user_id": original.user_id, "max_budget": 100, "alias": "Changed"} + bulk: Final = route.endswith("bulk_update") + payload: Final = {"customers": [update]} if bulk else update + response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == (200 if bulk else 404), response.text + if bulk: + assert response.json()["results"] == [ + {"user_id": original.user_id, "success": False, "updated_customer": None, "error": "Budget does not exist"} + ] + else: + assert response.json()["error"]["message"] == "Budget does not exist" + assert database.customers == {original.user_id: original} + assert database.budgets == {} + assert database.locked_budget_ids == {"gone"} + + +@pytest.mark.parametrize( + "route", ("/customer/update", "/end_user/update", "/customer/bulk_update", "/end_user/bulk_update") +) +@pytest.mark.parametrize("max_budget", (None, 100)) +@pytest.mark.parametrize("budget_id", (None, "other-budget")) +def test_customer_update_rejects_both_budget_fields_before_writing( + mock_prisma_client: MagicMock, + mock_user_api_key_auth: None, + route: str, + max_budget: int | None, + budget_id: str | None, +) -> None: + original: Final = LiteLLM_EndUserTable(user_id="customer", alias="Original", blocked=False) + database: Final = _CustomerBatchDatabase((original,)) + database.attach(mock_prisma_client) + invalid: Final = {"user_id": original.user_id, "max_budget": max_budget, "budget_id": budget_id} + payload: Final = ( + {"customers": [{"user_id": original.user_id, "alias": "Changed"}, invalid]} + if route.endswith("bulk_update") + else invalid + ) + response: Final = client.post(route, json=payload, headers={"Authorization": "Bearer test-key"}) + + assert response.status_code == 422, response.text + assert "Set either 'max_budget' or 'budget_id', not both." in response.text + assert database.customers == {original.user_id: original} + assert database.budgets == {} + mock_prisma_client.db.query_raw.assert_not_awaited() + mock_prisma_client.db.litellm_endusertable.update.assert_not_awaited() 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 e451c38a251..934eb9eac6b 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -1138,7 +1138,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"]}}, ) @@ -1182,7 +1182,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"]}}, ) @@ -1225,7 +1225,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 a5cae02ea10..2422fbc6e31 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 33557713162..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, @@ -1519,25 +1521,174 @@ _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 +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", 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/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 new file mode 100644 index 00000000000..e06ec7a9ae5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerForm.tsx @@ -0,0 +1,284 @@ +"use client"; + +import { useState } from "react"; +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 { cn } from "@/lib/cva.config"; +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" }, +]; + +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, baseline: EndUser | null) => void | Promise; + onCancel: () => void; +}; + +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 [baseline] = useState(customer); + const [initialValues] = useState(() => ({ + ...toCustomerFormValues(baseline), + ...(isEdit ? { budget_id: "" } : {}), + })); + const form = useZodForm(isEdit ? editCustomerSchema : customerSchema, { + defaultValues: initialValues, + }); + const hasCurrentBudget = isEdit && Boolean(customer?.budget_id); + const budgetFieldCopy = getBudgetFieldCopy(hasCurrentBudget, isEdit); + const handleSubmit = form.handleSubmit(async (values) => { + const isMaximumEdit = !values.budget_id && Boolean(values.max_budget.trim()); + if (hasCurrentBudget && hasSharedBudget !== false && isMaximumEdit) { + form.setError("max_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; + } + await onSubmit(values, baseline); + }); + const selectedBudgetId = form.watch("budget_id"); + const budgetOptions = [ + { value: "", label: budgetFieldCopy.emptyOption }, + ...(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..0175d0cc590 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerInfoView.tsx @@ -0,0 +1,289 @@ +"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 { getBudgetDurationLabel } from "@/components/common_components/budget_duration_dropdown"; +import MCPServerPermissions from "@/components/permissions/MCPServerPermissions"; +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 } 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; + hasSharedBudget: boolean | undefined; + 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 CustomerMcpAccess({ customer, accessToken }: Pick) { + const permission = customer.object_permission; + return ( + + ); +} + +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, + hasSharedBudget, + isEdit, + tab, + onTabChange, + onBack, + onEdit, + onCancelEdit, + ...actions +}: CustomerInfoViewProps) { + const { saveCustomer, isSaving } = useSaveCustomer(customer); + const handleSave = async (values: CustomerFormValues, baseline: EndUser | null) => { + try { + await saveCustomer(values, baseline); + 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..33f7a28d2df --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.integration.test.tsx @@ -0,0 +1,889 @@ +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]; + +const confirmBlocking = async (user: ReturnType, 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(); + 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("blocks selected customers in one batch and clears selection after success", async () => { + authorization.premiumUser = true; + const user = userEvent.setup(); + renderWithProviders(); + expect(screen.queryByRole("button", { name: "Block selected" })).not.toBeInTheDocument(); + 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", + body: { user_ids: customers.map((customer) => customer.user_id) }, + }), + ); + await waitFor(() => expect(screen.queryByRole("button", { name: "Block selected" })).not.toBeInTheDocument()); + expect(screen.getByRole("checkbox", { name: `Select customer ${activeCustomer.user_id}` })).not.toBeChecked(); + }); + + it.each(["Block", "Unblock"])("gates bulk %s with the Enterprise explanation", async (action) => { + const user = userEvent.setup(); + renderWithProviders(); + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + 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.", + ); + expect(button).toHaveAccessibleDescription( + "This is a LiteLLM Enterprise feature, and requires a valid key to use.", + ); + await user.keyboard("{Enter}"); + expect(post).not.toHaveBeenCalled(); + }); + + it("unblocks only selected IDs in one batch with Enterprise", async () => { + authorization.premiumUser = true; + const user = userEvent.setup(); + 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", + body: { user_ids: [blockedCustomer.user_id] }, + }), + ); + }); + + it("keeps selection and blocked state when bulk Unblock has no configured callback", async () => { + authorization.premiumUser = true; + const user = userEvent.setup(); + const message = "Blocked user check was never set"; + post.mockRejectedValue(new ApiError(message, 400, { detail: { error: message } })); + 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); + 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(); + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/unblock", { + accessToken: "sk-test", + body: { user_ids: [blockedCustomer.user_id] }, + }); + }); + + 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/ })).getByText("$0.00")).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(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); + 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 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(); + expect(within(bob).getByText("No restriction")).toBeInTheDocument(); + }); + + it.each([ + { surface: "row", action: "Block", customer: activeCustomer }, + { surface: "detail", action: "Block", customer: activeCustomer }, + { surface: "row", action: "Unblock", customer: blockedCustomer }, + { surface: "detail", action: "Unblock", customer: blockedCustomer }, + ])("disables $action without Enterprise in the $surface menu", async ({ surface, action, customer }) => { + const user = userEvent.setup(); + renderWithProviders( + , + surface === "detail" ? { searchParams: `customer=${encodeURIComponent(customer.user_id)}` } : undefined, + ); + await user.click( + screen.getByRole("button", { + name: surface === "detail" ? "More customer actions" : `Actions for ${customer.user_id}`, + }), + ); + 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.", + ); + expect(menuAction).toHaveAccessibleDescription( + "This is a LiteLLM Enterprise feature, and requires a valid key to use.", + ); + await user.keyboard("{Enter}"); + expect(post).not.toHaveBeenCalled(); + }); + + it.each([ + { surface: "row", action: "Block", endpoint: "/customer/block", customer: activeCustomer }, + { surface: "detail", action: "Block", endpoint: "/customer/block", customer: activeCustomer }, + { surface: "row", action: "Unblock", endpoint: "/customer/unblock", customer: blockedCustomer }, + { surface: "detail", action: "Unblock", endpoint: "/customer/unblock", customer: blockedCustomer }, + ])("uses $endpoint with Enterprise in the $surface menu", async ({ surface, action, endpoint, customer }) => { + authorization.premiumUser = true; + const user = userEvent.setup(); + renderWithProviders( + , + surface === "detail" ? { searchParams: `customer=${encodeURIComponent(customer.user_id)}` } : undefined, + ); + expect(screen.queryByRole("button", { name: action })).not.toBeInTheDocument(); + await user.click( + screen.getByRole("button", { + name: surface === "detail" ? "More customer actions" : `Actions for ${customer.user_id}`, + }), + ); + 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", + body: { user_ids: [customer.user_id] }, + }), + ); + }); + + 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.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(); + 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" }, + }), + ); + await waitFor(() => + expect(onUrlUpdate).toHaveBeenCalledWith( + expect.objectContaining({ + searchParams: expect.any(URLSearchParams), + }), + ), + ); + expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("edit")).toBeNull(); + }); + + 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.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [directCustomer] : 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" }, + }), + ); + }); + + it("denies every MCP server to a customer through the No MCP Servers option", async () => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + await user.click(await screen.findByLabelText("MCP Servers / Access Groups / Toolsets")); + const noServers = (await screen.findAllByRole("option")).find((option) => + option.textContent?.startsWith("No MCP Servers"), + ); + await user.click(noServers!); + await user.click(document.body); + 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, + object_permission: { + mcp_servers: ["no-mcp-servers"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: {}, + }, + }, + }), + ); + }); + + 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.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [directCustomer] : 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" }); + 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("checkbox", { name: "Select all rows" })).not.toBeInTheDocument(); + 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("uses the shared spend/budget cell and hides model details by default", () => { + 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 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("of $100.00")).toBeInTheDocument(); + expect(within(row).getByText("TPM: 600")).toBeInTheDocument(); + expect(within(row).getByText("RPM: 6")).toBeInTheDocument(); + expect(within(row).queryByText("All models the key can use")).not.toBeInTheDocument(); + }); + + it("shows Overview cards and shared permissions before entering Settings", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + const customer: EndUser = { + ...activeCustomer, + object_permission: { + ...activeCustomer.object_permission!, + vector_stores: ["vector-1"], + search_tools: ["search-1"], + skills: ["skill-1"], + }, + 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.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"); + 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("MCP Servers")).toBeInTheDocument(); + for (const unenforced of ["Vector Stores", "Search tools", "Skills", "vector-1", "search-1", "skill-1"]) { + expect(screen.queryByText(unenforced)).not.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(); + 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("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, 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" }; + 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"), "Keep current 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 shared 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 }, + }), + ); + }); + + it.each(["row", "detail"])( + "shows an Unblock callback error and keeps the customer blocked in the %s menu", + async (surface) => { + authorization.premiumUser = true; + const user = userEvent.setup(); + const message = "Blocked user check was never set"; + post.mockRejectedValue(new ApiError(message, 400, { detail: { error: message } })); + renderWithProviders( + <> + + + , + surface === "detail" ? { searchParams: "customer=bob%40example.test" } : undefined, + ); + await user.click( + screen.getByRole("button", { + name: surface === "detail" ? "More customer actions" : `Actions for ${blockedCustomer.user_id}`, + }), + ); + 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(); + expect(screen.queryByText("1 customer unblocked")).not.toBeInTheDocument(); + expect(screen.getByText("Blocked")).toBeInTheDocument(); + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/unblock", { + accessToken: "sk-test", + body: { user_ids: [blockedCustomer.user_id] }, + }); + }, + ); + it("keeps conflicting permission edits unsaved until Settings is reopened", async () => { + const refreshed: EndUser = { + ...activeCustomer, + object_permission: { + ...activeCustomer.object_permission!, + mcp_servers: ["server-1"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-1": ["read_document"] }, + }, + }; + const current = { customer: activeCustomer }; + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [current.customer] : current.customer, + isLoading: false, + isError: false, + })); + const user = userEvent.setup(); + const view = renderWithProviders( + <> + + + , + { + searchParams: "customer=alice%40example.test&edit=true", + }, + ); + const selectNoServers = async () => { + await user.click(await screen.findByLabelText("MCP Servers / Access Groups / Toolsets")); + await user.click(await screen.findByRole("option", { name: /No MCP Servers/ })); + await user.keyboard("{Escape}"); + }; + await selectNoServers(); + current.customer = refreshed; + view.rerender( + <> + + + , + ); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + expect(await screen.findByText(/Customer permissions or budget changed while editing/)).toBeInTheDocument(); + expect(post).not.toHaveBeenCalled(); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await user.click(screen.getByRole("button", { name: "Edit Settings" })); + await selectNoServers(); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { + user_id: refreshed.user_id, + object_permission: { + mcp_servers: ["no-mcp-servers"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: {}, + }, + }, + }), + ); + }); + + it("does not restore revoked permissions after an alias-only save and customer refresh", async () => { + const before: EndUser = { + ...activeCustomer, + models: ["old-model"], + allowed_model_region: "us", + default_model: "old-model", + budget_id: "old-budget", + object_permission: { + ...activeCustomer.object_permission!, + mcp_servers: ["server-1"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-1": ["read_document"] }, + }, + }; + const after: EndUser = { + ...before, + models: ["restricted-model"], + allowed_model_region: "eu", + default_model: "new-model", + budget_id: "new-budget", + object_permission: { ...before.object_permission!, mcp_servers: ["no-mcp-servers"], mcp_tool_permissions: {} }, + }; + let current = before; + useQuery.mockImplementation((_method: string, path: string) => ({ + data: path === "/customer/list" ? [current] : current, + isLoading: false, + isError: false, + })); + fetchServers.mockResolvedValue([{ server_id: "server-1", server_name: "Research Server" }]); + listTools.mockResolvedValue({ + error: false, + tools: [{ name: "read_document", description: "Read", inputSchema: {} }], + }); + const user = userEvent.setup(); + const view = renderWithProviders(, { searchParams: "customer=alice%40example.test&edit=true" }); + await screen.findByRole("radio", { name: "Flat List" }); + fireEvent.change(screen.getByLabelText("Alias"), { target: { value: "Updated Alice" } }); + current = after; + view.rerender(); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + await waitFor(() => expect(post).toHaveBeenCalledOnce()); + expect(post).toHaveBeenCalledExactlyOnceWith("/customer/update", { + accessToken: "sk-test", + body: { user_id: before.user_id, alias: "Updated Alice" }, + }); + }); +}); 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..0630c0e145d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersPanel.tsx @@ -0,0 +1,215 @@ +"use client"; + +import { Contact, Plus } from "lucide-react"; +import { parseAsBoolean, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; +import { useCallback, useMemo, useState } from "react"; + +import { + useCustomer, + useDeleteCustomers, + useSetCustomersBlockedState, +} 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 { CustomerBlockConfirmation } from "./CustomerBlockConfirmation"; +import { CreateCustomerModal } from "./CreateCustomerModal"; +import CustomerInfoView from "./CustomerInfoView"; +import CustomersTable from "./CustomersTable"; +import { getCustomerBudgetCounts, hasSharedCustomerBudget } from "./customerPayload"; + +const CUSTOMER_VIEW_PARSERS = { + customer: parseAsString, + edit: parseAsBoolean.withDefault(false), + customer_tab: parseAsStringLiteral(["overview", "settings"]).withDefault("overview"), +}; + +function CustomersHeader({ canModify, onCreate }: { canModify: boolean; onCreate: () => void }) { + return ( + + + + Customers + + Manage customer 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 [blockTarget, setBlockTarget] = useState(null); + const customers = useCustomers(); + const customer = useCustomer(customerId); + const budgetCounts = useMemo(() => getCustomerBudgetCounts(customers.data ?? []), [customers.data]); + const deleteCustomers = useDeleteCustomers(); + const setBlockedState = useSetCustomersBlockedState(); + const isBlocking = setBlockedState.isPending; + + 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 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); + } + }; + + const handleSetBlocked = async (userIds: string[], blocked: boolean): Promise => { + const canSetBlocked = canModify && !isBlocking && premiumUser; + if (!userIds.length || !canSetBlocked) return false; + try { + await setBlockedState.mutateAsync({ userIds, blocked }); + toast.success( + `${userIds.length} customer${userIds.length === 1 ? "" : "s"} ${blocked ? "blocked" : "unblocked"}`, + ); + return true; + } catch (error) { + toast.fromError(error); + return false; + } + }; + + 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={isBlocking} + onBack={closeCustomer} + onEdit={() => editCustomer(customerId)} + onCancelEdit={cancelEdit} + onDelete={setDeleteTarget} + onSetBlocked={setBlockTarget} + /> + ) : ( + <> + setIsCreateOpen(true)} /> + void customers.refetch()} + isRefreshing={customers.isFetching} + canModify={canModify} + premiumUser={premiumUser} + onView={openCustomer} + onEdit={editCustomer} + onDelete={setDeleteTarget} + onSetBlocked={setBlockTarget} + /> + + )} + {canModify && ( + <> + setBlockTarget(null)} + onConfirm={async () => { + if (blockTarget && (await handleSetBlocked([blockTarget.user_id], !blockTarget.blocked))) { + setBlockTarget(null); + } + }} + /> + + 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.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.integration.test.tsx new file mode 100644 index 00000000000..62f74e49e4b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.integration.test.tsx @@ -0,0 +1,222 @@ +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import Papa from "papaparse"; +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { toast } from "@/lib/toast"; +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { chooseSelectOption, renderWithProviders } from "../../../../../tests/test-utils"; +import CustomersTable from "./CustomersTable"; + +const customer: EndUser = { + user_id: "customer-01", + alias: "Research", + blocked: false, + spend: 0, + models: ["model-1", "model-2"], + litellm_budget_table: { max_budget: 0, created_at: "2026-01-01T00:00:00Z" }, + object_permission: { + object_permission_id: "permission-1", + mcp_servers: ["server-1"], + mcp_access_groups: ["group-1"], + mcp_toolsets: ["toolset-1"], + mcp_tool_permissions: {}, + agent_access_groups: [], + agents: [], + models: [], + blocked_tools: [], + search_tools: [], + vector_stores: [], + }, +}; +const props = { + data: [customer], + isLoading: false, + isError: false, + isRefreshing: false, + isBlocking: false, + canModify: true, + premiumUser: true, + onRefresh: vi.fn(), + onView: vi.fn(), + onEdit: vi.fn(), + onDelete: vi.fn(), + onSetBlocked: vi.fn(), + onBulkSetBlocked: vi.fn().mockResolvedValue(true), +}; + +afterEach(() => { + vi.clearAllMocks(); + vi.unstubAllGlobals(); + vi.restoreAllMocks(); +}); + +describe("CustomersTable controls", () => { + it("combines all three filters, writes URL state, and resets them", async () => { + const rows: EndUser[] = [ + customer, + { ...customer, user_id: "customer-02", alias: "Blocked", blocked: true }, + { ...customer, user_id: "customer-03", alias: "Other", object_permission: null, litellm_budget_table: null }, + ]; + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + await user.click(screen.getByRole("button", { name: "Filters" })); + expect(screen.getByRole("combobox", { name: "Budget" })).toHaveTextContent(/^All\b/); + await chooseSelectOption(user, screen.getByRole("combobox", { name: "Status" }), "Active"); + await chooseSelectOption(user, screen.getByRole("combobox", { name: "Budget" }), "Has budget"); + await chooseSelectOption(user, screen.getByRole("combobox", { name: "MCP access" }), "Limited MCP access"); + expect(screen.getByRole("combobox", { name: "Budget" })).toHaveTextContent(/^Has budget\b/); + await user.click(screen.getByRole("button", { name: "Apply Filters" })); + + await waitFor(() => expect(screen.queryByRole("row", { name: /customer-02/ })).not.toBeInTheDocument()); + expect(screen.queryByRole("row", { name: /customer-03/ })).not.toBeInTheDocument(); + expect(screen.getByRole("row", { name: /customer-01/ })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Unblock selected" })).not.toBeInTheDocument(); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + const params = onUrlUpdate.mock.calls.at(-1)?.[0].searchParams; + expect(params.get("filter_blocked")).toBe("active"); + expect(params.get("filter_max_budget")).toBe("yes"); + expect(params.get("filter_mcp_access")).toBe("limited"); + await user.click(screen.getByRole("button", { name: /^Filters/ })); + await user.click(screen.getByRole("button", { name: "Reset" })); + await user.click(screen.getByRole("button", { name: "Apply Filters" })); + expect(await screen.findByRole("row", { name: /customer-02/ })).toBeInTheDocument(); + expect(screen.getByRole("row", { name: /customer-03/ })).toBeInTheDocument(); + }); + + it("restores negative filters from the URL and clears selection on search changes", async () => { + const blocked: EndUser = { + ...customer, + user_id: "customer-02", + blocked: true, + object_permission: null, + litellm_budget_table: null, + }; + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { + searchParams: "filter_blocked=blocked&filter_max_budget=no&filter_mcp_access=no", + onUrlUpdate, + }); + expect(screen.getByRole("row", { name: /customer-02/ })).toBeInTheDocument(); + expect(screen.queryByRole("row", { name: /customer-01/ })).not.toBeInTheDocument(); + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + fireEvent.change(screen.getByRole("textbox", { name: "Search customers" }), { target: { value: "missing" } }); + await waitFor(() => expect(screen.queryByRole("row", { name: /customer-02/ })).not.toBeInTheDocument()); + expect(screen.queryByRole("button", { name: "Unblock selected" })).not.toBeInTheDocument(); + await waitFor(() => + expect(onUrlUpdate).toHaveBeenLastCalledWith( + expect.objectContaining({ + searchParams: expect.any(URLSearchParams), + queryString: expect.stringContaining("search=missing"), + }), + ), + ); + fireEvent.change(screen.getByRole("textbox", { name: "Search customers" }), { target: { value: "" } }); + expect(await screen.findByRole("checkbox", { name: "Select customer customer-02" })).not.toBeChecked(); + }); + + it("selects the current page and retains IDs selected on other pages", async () => { + const rows = Array.from( + { length: 26 }, + (_, index): EndUser => ({ + ...customer, + user_id: `customer-${String(index).padStart(2, "0")}`, + alias: null, + }), + ); + const user = userEvent.setup(); + renderWithProviders(, { searchParams: "page_size=25" }); + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + expect(screen.getByText("25 selected")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Go to next page" })); + expect(await screen.findByRole("checkbox", { name: "Select customer customer-25" })).not.toBeChecked(); + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + expect(screen.getByText("26 selected")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Unblock selected" })); + const dialog = await screen.findByRole("dialog", { name: "Unblock 26 customers?" }); + expect(props.onBulkSetBlocked).not.toHaveBeenCalled(); + await user.click(within(dialog).getByRole("button", { name: "Unblock" })); + expect(props.onBulkSetBlocked).toHaveBeenCalledExactlyOnceWith( + rows.map((row) => row.user_id), + false, + ); + await waitFor(() => expect(screen.queryByRole("button", { name: "Unblock selected" })).not.toBeInTheDocument()); + }); + + it.each([ + { blocked: false, action: "Block" }, + { blocked: true, action: "Unblock" }, + ])("disables selection, row $action, and batch controls during a block mutation", async ({ blocked, action }) => { + const data = [{ ...customer, blocked }]; + const user = userEvent.setup(); + const view = renderWithProviders(); + await user.click(screen.getByRole("checkbox", { name: "Select all rows" })); + view.rerender(); + expect(screen.getByRole("button", { name: "Block selected" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Unblock selected" })).toBeDisabled(); + expect(screen.getByRole("button", { name: "Clear selection" })).toBeDisabled(); + await user.click(screen.getByRole("checkbox", { name: "Select customer customer-01" })); + expect(screen.getByRole("checkbox", { name: "Select customer customer-01" })).toBeChecked(); + await user.click(screen.getByRole("button", { name: "Actions for customer-01" })); + expect(await screen.findByRole("menuitem", { name: action })).toHaveAttribute("aria-disabled", "true"); + expect(props.onBulkSetBlocked).not.toHaveBeenCalled(); + expect(props.onSetBlocked).not.toHaveBeenCalled(); + }); + + it.each([ + { search: "matching", count: 26 }, + { search: "customer-00", count: 1 }, + ])("exports $count matching customers using search and filters", async ({ search, count }) => { + const rows = Array.from( + { length: 28 }, + (_, index): EndUser => ({ + ...customer, + user_id: `customer-${String(index).padStart(2, "0")}`, + alias: index === 27 ? "Other" : "Matching", + blocked: index === 26, + }), + ); + const createObjectURL = vi.fn().mockReturnValue("blob:customers-export"); + const revokeObjectURL = vi.fn(); + vi.stubGlobal( + "URL", + class extends URL { + static createObjectURL = createObjectURL; + static revokeObjectURL = revokeObjectURL; + }, + ); + const download = vi.spyOn(HTMLAnchorElement.prototype, "click").mockImplementation(() => {}); + const user = userEvent.setup(); + renderWithProviders(, { + searchParams: `page_size=25&search=${search}&filter_blocked=active&filter_mcp_access=yes`, + }); + expect(screen.getByText(`Showing 1-${Math.min(count, 25)} of ${count}`)).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Export CSV" })); + const blob = createObjectURL.mock.calls[0][0] as Blob; + const csv = await new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(String(reader.result)); + reader.onerror = () => reject(reader.error); + reader.readAsText(blob); + }); + const parsed = Papa.parse>(csv, { header: true }); + expect(parsed.errors).toEqual([]); + expect(parsed.data.map((row) => row["Customer ID"])).toEqual(rows.slice(0, count).map((row) => row.user_id)); + expect(download).toHaveBeenCalledOnce(); + expect(toast.success).toHaveBeenCalledWith(`Exported ${count} customer${count === 1 ? "" : "s"}`); + expect(revokeObjectURL).toHaveBeenCalledExactlyOnceWith("blob:customers-export"); + }); + + it("wires the Resources badges into the customer table", () => { + renderWithProviders(); + expect(screen.getByRole("columnheader", { name: "Resources" })).toBeInTheDocument(); + const row = screen.getByRole("row", { name: /customer-01/ }); + expect(within(row).getByLabelText("2 models")).toHaveTextContent("2"); + expect(within(row).getByLabelText("1 MCP server")).toHaveTextContent("1"); + expect(within(row).getByLabelText("1 MCP access group")).toHaveTextContent("1"); + expect(within(row).getByLabelText("1 MCP toolset")).toHaveTextContent("1"); + }); +}); 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..f262c14797b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTable.tsx @@ -0,0 +1,296 @@ +"use client"; + +import type { RowSelectionState } from "@tanstack/react-table"; +import { Ban, Contact, Download, Search, ShieldCheck } from "lucide-react"; +import { useId, useMemo, useState } from "react"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { + createSelectionColumn, + DataTable, + DataTableFilterDrawer, + DataTableFilterField, + DataTableToolbar, + useUrlTableState, +} from "@/components/shared/DataTable"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; +import { toast } from "@/lib/toast"; + +import { CustomerBlockConfirmation } from "./CustomerBlockConfirmation"; +import { exportCustomersToCsv } from "./customerCsvExport"; +import { + CUSTOMER_FILTER_FIELDS, + customerFilterFn, + customerMatchesSearch, + formatCustomerFilterValue, +} from "./customerFilters"; +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: CUSTOMER_FILTER_FIELDS.map((field) => field.id), +}; + +const DEFAULT_COLUMN_VISIBILITY = { mcp_access: false, models: false, max_budget: false }; +const FILTER_LABELS = Object.fromEntries(CUSTOMER_FILTER_FIELDS.map((field) => [field.id, field.label])); + +type CustomersTableProps = CustomerTableActions & { + data: EndUser[]; + isLoading: boolean; + isError: boolean; + onRefresh: () => void; + isRefreshing: boolean; + onBulkSetBlocked: (userIds: string[], blocked: boolean) => Promise; +}; + +function EmptyCustomers() { + return ( +
+ +

No customers found

+

Create a customer to manage customer limits and access

+
+ ); +} + +function BulkBlockedStateButton({ + blocked, + premiumUser, + disabled, + onClick, +}: { + blocked: boolean; + premiumUser: boolean; + disabled: boolean; + onClick: () => void; +}) { + const descriptionId = useId(); + const content = ( + <> + {blocked ? : } + {blocked ? "Block selected" : "Unblock selected"} + + ); + if (premiumUser) + return ( + + ); + return ( + + + } + > + {content} + + + This is a LiteLLM Enterprise feature, and requires a valid key to use. + + + + ); +} + +export default function CustomersTable({ + data, + isLoading, + isError, + onRefresh, + isRefreshing, + isBlocking, + onBulkSetBlocked, + ...actions +}: CustomersTableProps) { + const { + search, + setSearch, + sorting, + onSortingChange, + pagination, + onPaginationChange, + columnFilters, + 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); + if (selectionScope !== filterScope) { + setSelectionScope(filterScope); + setRowSelection({}); + } + const { canModify, premiumUser, onView, onEdit, onDelete, onSetBlocked } = actions; + const columns = useMemo(() => { + const columnActions = { canModify, premiumUser, isBlocking, onView, onEdit, onDelete, onSetBlocked }; + const customerColumns = getCustomerTableColumns(columnActions).map((column) => + CUSTOMER_FILTER_FIELDS.some((field) => field.id === column.id) + ? { ...column, filterFn: customerFilterFn } + : column, + ); + return canModify + ? [ + createSelectionColumn({ rowAriaLabel: (row) => `Select customer ${row.original.user_id}` }), + ...customerColumns, + ] + : customerColumns; + }, [canModify, premiumUser, isBlocking, onView, onEdit, onDelete, onSetBlocked]); + const searchedCustomers = useMemo( + () => data.filter((customer) => customerMatchesSearch(customer, search)), + [data, search], + ); + + const handleBulkSetBlocked = async (userIds: string[], blocked: boolean) => { + if (await onBulkSetBlocked(userIds, blocked)) { + setRowSelection({}); + setBlockTarget(null); + } + }; + + return ( +
+ setBlockTarget(null)} + onConfirm={() => { + if (blockTarget) void handleBulkSetBlocked(blockTarget.userIds, blockTarget.blocked); + }} + /> + customer.user_id} + defaultColumnVisibility={DEFAULT_COLUMN_VISIBILITY} + enableRowSelection={canModify && !isBlocking} + rowSelection={rowSelection} + onRowSelectionChange={(updater) => { + if (!isBlocking) setRowSelection(updater); + }} + toolbar={(table) => { + const selectedIds = table.getFilteredSelectedRowModel().rows.map((row) => row.original.user_id); + return ( + <> +
+
+ + setSearch(event.target.value)} + /> +
+ setFiltersOpen(true)} + filterLabels={FILTER_LABELS} + formatFilterValue={formatCustomerFilterValue} + > + + +
+ {canModify && selectedIds.length > 0 && ( +
+ {selectedIds.length} selected + setBlockTarget({ userIds: selectedIds, blocked: true })} + /> + setBlockTarget({ userIds: selectedIds, blocked: false })} + /> + +
+ )} + + {({ get, set }) => + CUSTOMER_FILTER_FIELDS.map((field) => ( + + + + )) + } + + + ); + }} + filterMode="client" + columnFilters={columnFilters} + onColumnFiltersChange={onColumnFiltersChange} + 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.tsx b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.tsx new file mode 100644 index 00000000000..96e2d23bbc0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.test.tsx @@ -0,0 +1,105 @@ +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, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; + +import { CustomerResourcesCell, describeMcpGrant } 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"); + }); +}); + +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 server, 1 tool limit"); + }); + + it("does not describe the unsupported customer all-servers sentinel as granting access", () => { + const permission = { ...emptyGrant, mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL] }; + expect(describeMcpGrant(permission)).toBe("No MCP access"); + }); + + 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..91edd53fadc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomersTableColumns.tsx @@ -0,0 +1,320 @@ +"use client"; + +import { useId } from "react"; + +import type { ColumnDef } from "@tanstack/react-table"; +import { Ban, Copy, Layers, MoreHorizontal, Pencil, Server, ShieldCheck, Trash2, Network, 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"; +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { IdentityCell, ModelsCell, SpendBudgetCell, 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"; + +import { customerMcpScope, getCustomerResourceCounts } from "./customerResources"; + +export type CustomerTableActions = { + canModify: boolean; + premiumUser: boolean; + isBlocking: 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 descriptionId = useId(); + const label = customer.blocked ? "Unblock" : "Block"; + const content = ( + <> + {customer.blocked ? : } + {label} + + ); + if (premiumUser) { + return ( + onSetBlocked(customer)} + > + {content} + + ); + } + return ( + + + } + > + {content} + + + 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 scope = customerMcpScope(permission); + if (scope === "denied") return "No MCP access"; + if (scope === "unrestricted") return "No restriction"; + const counts = getCustomerResourceCounts({ object_permission: permission }); + const toolPermissionKeys = Object.keys(permission?.mcp_tool_permissions ?? {}).filter( + (server) => server !== ALL_PROXY_MCP_SERVERS_SENTINEL && server !== NO_MCP_SERVERS_SENTINEL, + ); + const parts = [ + countLabel(counts.servers, "server"), + countLabel(counts.accessGroups, "group"), + countLabel(counts.toolsets, "toolset"), + countLabel(toolPermissionKeys.length, "tool limit"), + ].filter(Boolean); + return parts.join(", "); +} + +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-violet-50 text-violet-700 ring-violet-600/20 dark:bg-violet-950 dark:text-violet-300 dark:ring-violet-400/30", + }, + { + key: "accessGroups", + label: "MCP access groups", + icon: Network, + className: "bg-cyan-50 text-cyan-700 ring-cyan-600/20 dark:bg-cyan-950 dark:text-cyan-300 dark:ring-cyan-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, + premiumUser, + onEdit, + onDelete, + onSetBlocked, + isBlocking, +}: 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: "resources", + meta: { title: "Resources", skeleton: "chips" }, + header: "Resources", + size: 250, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "spend", + accessorKey: "spend", + meta: { title: "Spend / Budget", skeleton: "meter", numeric: true }, + header: ({ column }) => , + size: 200, + cell: ({ row }) => ( + + ), + }, + { + id: "max_budget", + accessorFn: (customer) => customer.litellm_budget_table?.max_budget, + enableHiding: false, + enableSorting: false, + }, + { + 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/customerCsvExport.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.test.ts new file mode 100644 index 00000000000..5201baab224 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.test.ts @@ -0,0 +1,172 @@ +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"], + mcp_tool_permissions: { "server-9": ["write", "read"], "server-1": ["list"] }, + 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", + "MCP Tool Permissions": '{"server-1":["list"],"server-9":["read","write"]}', + "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", "=SUM(A1:A9)\nsecond line", "+cmd\r\nsecond line"])( + "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..2e45200fe3a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerCsvExport.ts @@ -0,0 +1,55 @@ +import Papa from "papaparse"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +const cell = (value: string | number | null | undefined): string | number => value ?? ""; + +const toolPermissionsCell = (permissions: Readonly>): string => { + const servers = Object.keys(permissions).toSorted(); + if (servers.length === 0) return ""; + return JSON.stringify(Object.fromEntries(servers.map((server) => [server, permissions[server].toSorted()]))); +}; + +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(", "), + "MCP Tool Permissions": toolPermissionsCell(permission?.mcp_tool_permissions ?? {}), + "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: /^[=+\-@\t\r]/ }); + +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..727d5d352d8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.test.ts @@ -0,0 +1,90 @@ +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: [], +}; + +const MCP_SCOPES = ["limited", "denied", "unrestricted"] as const; + +const scopesMatching = (row: EndUser): string[] => + MCP_SCOPES.filter((scope) => customerMatchesFilter(row, "mcp_access", scope)); + +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(scopesMatching(row)).toEqual(["limited"]); + }); + + it.each([ + { mcp_servers: [NO_MCP_SERVERS_SENTINEL], mcp_access_groups: ["group-1"] }, + { mcp_servers: [ALL_PROXY_MCP_SERVERS_SENTINEL] }, + ])("matches denied MCP grants only as No MCP access %j", (grant) => { + expect(scopesMatching({ ...customer, object_permission: { ...permission, ...grant } })).toEqual(["denied"]); + }); + + it.each([undefined, null, permission])( + "matches customers without a grant only as No MCP restriction (%j)", + (object_permission) => { + expect(scopesMatching({ ...customer, object_permission })).toEqual(["unrestricted"]); + }, + ); + + 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, "mcp_access", "yes")).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", "unrestricted")).toBe("No MCP restriction"); + }); +}); 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..c078eeed42e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerFilters.ts @@ -0,0 +1,58 @@ +import type { FilterFn } from "@tanstack/react-table"; + +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; + +import { customerMcpScope } 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: "limited", label: "Limited MCP access" }, + { value: "denied", label: "No MCP access" }, + { value: "unrestricted", label: "No MCP restriction" }, + ], + }, +] 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 (columnId === "mcp_access") { + if (value !== "limited" && value !== "denied" && value !== "unrestricted") return true; + return customerMcpScope(customer.object_permission) === value; + } + if (value !== "yes" && value !== "no") return true; + if (columnId === "max_budget") return (customer.litellm_budget_table?.max_budget != null) === (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/customerPayload.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts new file mode 100644 index 00000000000..36e44343e9a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.test.ts @@ -0,0 +1,511 @@ +import { afterEach, 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 { NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { + buildCustomerPayload, + getCustomerBudgetCounts, + hasSharedCustomerBudget, + 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("shared customer budgets", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("counts customers per budget on browsers without Map.groupBy", () => { + vi.stubGlobal("Map", Object.assign(class extends Map {}, { groupBy: undefined })); + const linked = { ...customer, budget_id: "shared-budget" }; + + expect( + getCustomerBudgetCounts([ + linked, + { ...linked, user_id: "another-customer" }, + { ...customer, user_id: "no-budget" }, + ]), + ).toEqual(new Map([["shared-budget", 2]])); + }); + + it("allows a budget linked to only one customer", () => { + const owned = { ...customer, budget_id: "own-budget" }; + expect( + hasSharedCustomerBudget( + owned, + getCustomerBudgetCounts([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, getCustomerBudgetCounts([shared, { ...shared, user_id: "another-customer" }])), + ).toBe(true); + }); + + it("does not treat customers without budgets as sharing a budget", () => { + expect( + hasSharedCustomerBudget( + customer, + getCustomerBudgetCounts([customer, { ...customer, user_id: "another-customer" }]), + ), + ).toBe(false); + }); +}); + +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("does not revive stale tool grants when replacing No MCP Servers", () => { + const anotherServer = { ...server, server_id: "another-server" }; + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [NO_MCP_SERVERS_SENTINEL], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { [anotherServer.server_id]: ["read"] }, + }, + }; + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [server.server_id], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { ...stored.object_permission.mcp_tool_permissions, [server.server_id]: ["read"] }, + }; + + const expected = { + mcp_servers: [server.server_id], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { [server.server_id]: ["read"] }, + }; + expect(buildCustomerPayload(values, stored, [server, anotherServer], []).object_permission).toEqual(expected); + }); + + 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("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, + }); + }); + + it("omits unchanged model, region, default and budget settings on an alias edit", () => { + const stored: EndUser = { + ...customer, + models: ["model-a", "model-b"], + allowed_model_region: "eu", + default_model: "model-a", + litellm_budget_table: { max_budget: 50, created_at: "2026-01-01T00:00:00Z" }, + }; + const values = { ...toCustomerFormValues(stored), alias: "Updated alias", models: ["model-b", "model-a"] }; + + expect(buildCustomerPayload(values, stored, [], [])).toEqual({ + user_id: stored.user_id, + alias: values.alias, + }); + }); + + it("includes intentional model, region, default and budget changes", () => { + const stored: EndUser = { + ...customer, + models: ["model-a"], + allowed_model_region: "eu", + default_model: "model-a", + litellm_budget_table: { max_budget: 50, created_at: "2026-01-01T00:00:00Z" }, + }; + const values = { + ...toCustomerFormValues(stored), + models: [], + allowed_model_region: "us" as const, + default_model: "model-b", + max_budget: "25", + }; + + const expected = { + user_id: stored.user_id, + models: [], + allowed_model_region: "us", + default_model: "model-b", + max_budget: 25, + }; + expect(buildCustomerPayload(values, stored, [], [])).toEqual(expected); + }); + + 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: "", + }); + }); + + 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, + }); + }); + + 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, + }); + }); + + 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, + 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, + 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, + object_permission: grant, + }); + }); + + it.each([ + { edit: "only the alias", servers: ["server-a"], expected: undefined }, + { + edit: "the server selection", + servers: ["server-a", "server-c"], + expected: { + mcp_servers: ["server-a", "server-c"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "server-a": ["read"], "server-b": ["read"] }, + }, + }, + ])("keeps the other server behind a shared alias key when editing $edit", ({ servers, expected }) => { + const sharedServer = { ...server, server_id: "server-a", alias: "shared" }; + const anotherSharedServer = { ...server, server_id: "server-b", alias: "shared" }; + const addedServer = { ...server, server_id: "server-c" }; + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [sharedServer.server_id], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { shared: ["read"] }, + }, + }; + const values = { + ...toCustomerFormValues(stored), + alias: "New alias", + mcp_servers_and_groups: { servers, accessGroups: [], toolsets: [] }, + }; + + expect( + buildCustomerPayload(values, stored, [sharedServer, anotherSharedServer, addedServer], []).object_permission, + ).toEqual(expected); + }); + + it("keeps the other shared-alias server once the selected one is deselected", () => { + const sharedServer = { ...server, server_id: "server-a", alias: "shared" }; + const anotherSharedServer = { ...server, server_id: "server-b", alias: "shared" }; + const keptServer = { ...server, server_id: "server-c" }; + const stored = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [sharedServer.server_id, keptServer.server_id], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { shared: ["read"] }, + }, + }; + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [keptServer.server_id], accessGroups: [], toolsets: [] }, + }; + const grant = { + mcp_servers: [keptServer.server_id], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { [anotherSharedServer.server_id]: ["read"] }, + }; + + expect( + buildCustomerPayload(values, stored, [sharedServer, anotherSharedServer, keptServer], []).object_permission, + ).toEqual(grant); + }); + + it("keeps the stored ID of an existing customer exactly as stored", () => { + const stored = { ...customer, user_id: " customer-id " }; + + expect(buildCustomerPayload(toCustomerFormValues(stored), stored, [], []).user_id).toBe(" customer-id "); + }); + + 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, + 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..00fbeac2eda --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts @@ -0,0 +1,178 @@ +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 { NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { + accessGroupNamesOf, + mcpAllowedToolsFor, + mcpServersForIdentifier, +} from "@/components/mcp_server_management/effectiveMcpServers"; +import { + extractMcpEntitlementUpdate, + type McpEntitlementUpdate, +} from "@/components/mcp_server_management/mcpEntitlement"; + +export const getCustomerBudgetCounts = (customers: readonly EndUser[]): ReadonlyMap => + customers.reduce( + (counts, { budget_id }) => (budget_id ? counts.set(budget_id, (counts.get(budget_id) ?? 0) + 1) : counts), + new Map(), + ); + +export const hasSharedCustomerBudget = (customer: EndUser, budgetCounts: ReadonlyMap): boolean => + Boolean(customer.budget_id && (budgetCounts.get(customer.budget_id) ?? 0) > 1); + +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 stringsIdentity = (values: readonly string[]): string => JSON.stringify([...new Set(values)].sort()); + +const optionalCustomerFields = ( + values: CustomerFormValues, + existing: EndUser | null, +): Partial => ({ + ...(values.alias.trim() !== (existing?.alias ?? "") ? { alias: values.alias.trim() } : {}), + ...(values.budget_id && values.budget_id !== existing?.budget_id ? { budget_id: values.budget_id } : {}), + ...(values.max_budget.trim() && Number(values.max_budget) !== existing?.litellm_budget_table?.max_budget + ? { max_budget: Number(values.max_budget) } + : {}), + ...(values.allowed_model_region && values.allowed_model_region !== existing?.allowed_model_region + ? { allowed_model_region: values.allowed_model_region } + : {}), + ...(values.default_model.trim() && values.default_model.trim() !== existing?.default_model?.trim() + ? { default_model: values.default_model.trim() } + : {}), +}); + +export const buildCustomerPayload = ( + values: CustomerFormValues, + existing: EndUser | null, + servers: MCPServer[], + toolsets: MCPToolset[], +): UpdateCustomerRequest => { + const grant = extractMcpEntitlementUpdate(values, existing?.object_permission, servers, toolsets); + return { + user_id: existing?.user_id ?? values.user_id.trim(), + ...optionalCustomerFields(values, existing), + ...(stringsIdentity(values.models) !== stringsIdentity(existing?.models ?? []) ? { models: values.models } : {}), + ...(grant ? { object_permission: grant } : {}), + }; +}; + +const customerSecurityIdentity = (customer: EndUser): string => { + const grant = storedMcpGrant(customer); + const identity = { + models: stringsIdentity(customer.models ?? []), + allowed_model_region: customer.allowed_model_region ?? "", + default_model: customer.default_model ?? "", + budget_id: customer.budget_id ?? "", + max_budget: customer.litellm_budget_table?.max_budget ?? null, + mcp_servers: stringsIdentity(grant.mcp_servers), + mcp_access_groups: stringsIdentity(grant.mcp_access_groups), + mcp_toolsets: stringsIdentity(grant.mcp_toolsets), + mcp_tool_permissions: Object.keys(grant.mcp_tool_permissions) + .sort() + .map((key) => [key, stringsIdentity(grant.mcp_tool_permissions[key])]), + }; + return JSON.stringify(identity); +}; + +export const hasCustomerEditConflict = ( + payload: UpdateCustomerRequest, + baseline: EndUser | null, + current: EndUser | null, +): boolean => { + if (!baseline || !current) return baseline !== current; + if (baseline.user_id !== current.user_id) return true; + const changesSecurity = Object.keys(payload).some((key) => key !== "user_id" && key !== "alias"); + return changesSecurity && customerSecurityIdentity(baseline) !== customerSecurityIdentity(current); +}; + +export const directAndGroupServerIds = (grant: McpEntitlementUpdate, servers: MCPServer[]): ReadonlySet => { + if (grant.mcp_servers.includes(NO_MCP_SERVERS_SENTINEL)) return new Set(); + const directIds = grant.mcp_servers.flatMap((identifier) => + mcpServersForIdentifier(servers, identifier).map((server) => server.server_id), + ); + const groupIds = servers + .filter((server) => accessGroupNamesOf(server).some((group) => grant.mcp_access_groups.includes(group))) + .map((server) => server.server_id); + return new Set([...directIds, ...groupIds]); +}; + +const unrestrictedServerIds = ( + grant: McpEntitlementUpdate, + servers: MCPServer[], + toolsets: MCPToolset[], +): ReadonlySet => { + const toolsetServerIds = new Set( + toolsets + .filter((toolset) => grant.mcp_toolsets.includes(toolset.toolset_id)) + .flatMap((toolset) => toolset.tools.map((tool) => tool.server_id)), + ); + const grantedIds = directAndGroupServerIds(grant, servers); + return new Set( + servers + .filter((server) => grantedIds.has(server.server_id) && !toolsetServerIds.has(server.server_id)) + .filter((server) => mcpAllowedToolsFor(server, grant.mcp_tool_permissions, servers) === undefined) + .map((server) => server.server_id), + ); +}; + +export const serversLosingToolLimits = ( + values: CustomerFormValues, + existing: EndUser | null, + servers: MCPServer[], + toolsets: MCPToolset[], +): readonly MCPServer[] => { + const grant = extractMcpEntitlementUpdate(values, existing?.object_permission, servers, toolsets); + if (!grant) return []; + const stored = storedMcpGrant(existing); + const retainedIds = directAndGroupServerIds(stored, servers); + const directIds = new Set( + grant.mcp_servers.flatMap((identifier) => + mcpServersForIdentifier(servers, identifier).map((server) => server.server_id), + ), + ); + const unrestricted = unrestrictedServerIds(grant, servers, toolsets); + const alreadyUnrestricted = unrestrictedServerIds(stored, servers, toolsets); + const needsToolLimit = (server: MCPServer): boolean => { + const protectedGrant = retainedIds.has(server.server_id) || directIds.has(server.server_id); + return protectedGrant && unrestricted.has(server.server_id) && !alreadyUnrestricted.has(server.server_id); + }; + return servers.filter(needsToolLimit); +}; 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..25046fcc0ea --- /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 { customerMcpScope, 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(customerMcpScope(configured.object_permission)).toBe("limited"); + }); + + 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(customerMcpScope(configured.object_permission)).toBe("limited"); + }); + + 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(customerMcpScope(configured.object_permission)).toBe("denied"); + }); + + 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(customerMcpScope(configured.object_permission)).toBe("denied"); + }); + + 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(customerMcpScope(configured.object_permission)).toBe("unrestricted"); + }); +}); + +describe("customerMcpScope", () => { + it.each([ + undefined, + null, + emptyPermission, + { ...emptyPermission, mcp_servers: null, mcp_access_groups: null, mcp_toolsets: null }, + ])("treats an empty customer-level grant as unrestricted (%j)", (object_permission) => { + const unconfigured = { ...customer, object_permission }; + expect(customerMcpScope(object_permission)).toBe("unrestricted"); + 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(customerMcpScope({ ...emptyPermission, ...grant })).toBe("limited"); + }); +}); 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..49006149b15 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerResources.ts @@ -0,0 +1,44 @@ +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: Pick, +): 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 type CustomerMcpScope = "limited" | "denied" | "unrestricted"; + +export function customerMcpScope(permission: EndUser["object_permission"]): CustomerMcpScope { + const servers = permission?.mcp_servers ?? []; + if (servers.includes(NO_MCP_SERVERS_SENTINEL)) return "denied"; + const counts = getCustomerResourceCounts({ object_permission: permission }); + if (counts.servers > 0 || counts.accessGroups > 0 || counts.toolsets > 0) return "limited"; + const onlyUnsupportedAll = + servers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL) || + Object.hasOwn(permission?.mcp_tool_permissions ?? {}, ALL_PROXY_MCP_SERVERS_SENTINEL); + return onlyUnsupportedAll ? "denied" : "unrestricted"; +} 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..6b27ee5ac8a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.test.tsx @@ -0,0 +1,440 @@ +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(), + accessGroups: 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/mcpServers/useMCPAccessGroups", () => ({ useMCPAccessGroups: mocks.accessGroups })); +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.accessGroups.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" }); + }); + + it.each([ + { + field: "MCP grant", + current: { + ...customer, + object_permission: { ...permission, mcp_servers: ["no-mcp-servers"], mcp_tool_permissions: {} }, + }, + }, + { field: "models", current: { ...customer, models: ["model-a"] } }, + { field: "region", current: { ...customer, allowed_model_region: "eu" as const } }, + { field: "default model", current: { ...customer, default_model: "model-a" } }, + { field: "budget link", current: { ...customer, budget_id: "new-budget" } }, + { + field: "maximum budget", + current: { ...customer, litellm_budget_table: { max_budget: 25, created_at: "2026-01-01T00:00:00Z" } }, + }, + ])("refuses a deliberate permission edit after the stored $field changes", async ({ current }) => { + mocks.servers.mockReturnValue({ data: [server], isSuccess: true }); + const { result, rerender } = renderHook(({ current }) => useSaveCustomer(current), { + initialProps: { current: customer }, + }); + const values = { + ...toCustomerFormValues(customer), + mcp_tool_permissions: { [server.server_id]: ["read", "write"] }, + }; + rerender({ current }); + + await expect(result.current.saveCustomer(values, customer)).rejects.toThrow("changed while editing"); + expect(mocks.update).not.toHaveBeenCalled(); + }); + + it("allows deliberate permission edits after an unrelated refresh or reordered grants", async () => { + mocks.servers.mockReturnValue({ data: [server], isSuccess: true }); + const baseline: EndUser = { + ...customer, + models: ["model-a", "model-b"], + object_permission: { ...permission, mcp_tool_permissions: { [server.server_id]: ["read", "write"] } }, + }; + const current: EndUser = { + ...baseline, + alias: "Changed elsewhere", + spend: 10, + models: ["model-b", "model-a"], + object_permission: { ...permission, mcp_tool_permissions: { [server.server_id]: ["write", "read"] } }, + }; + const { result } = renderHook(() => useSaveCustomer(current)); + + await result.current.saveCustomer( + { + ...toCustomerFormValues(baseline), + mcp_tool_permissions: { [server.server_id]: ["read"] }, + }, + baseline, + ); + + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ + user_id: customer.user_id, + object_permission: { + ...clearedGrant, + mcp_servers: [server.server_id], + mcp_tool_permissions: { [server.server_id]: ["read"] }, + }, + }); + }); + + 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, + 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, + object_permission: clearedGrant, + }); + }); + + it("refuses a new direct server until it has an allowlist, then saves the allowlist", async () => { + const added: MCPServer = { ...server, server_id: "added-server" }; + mocks.servers.mockReturnValue({ data: [server, added], isSuccess: true }); + const { result } = renderHook(() => useSaveCustomer(customer)); + const values = { + ...toCustomerFormValues(customer), + mcp_servers_and_groups: { servers: [server.server_id, added.server_id], accessGroups: [], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow( + "Choose allowed tools for each MCP server before saving", + ); + expect(mocks.update).not.toHaveBeenCalled(); + + await result.current.saveCustomer({ + ...values, + mcp_tool_permissions: { ...values.mcp_tool_permissions, [added.server_id]: ["read"] }, + }); + + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ + user_id: customer.user_id, + object_permission: { + mcp_servers: [server.server_id, added.server_id], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { [server.server_id]: ["read"], [added.server_id]: ["read"] }, + }, + }); + }); + + it("refuses removing the toolset that limits a stored direct server's tools until it has an allowlist", async () => { + const stored: EndUser = { + ...customer, + object_permission: { ...permission, mcp_toolsets: [toolset.toolset_id], mcp_tool_permissions: {} }, + }; + mocks.servers.mockReturnValue({ data: [server], isSuccess: true }); + mocks.toolsets.mockReturnValue({ data: [toolset], isSuccess: true }); + const { result } = renderHook(() => useSaveCustomer(stored)); + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [server.server_id], accessGroups: [], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow( + "Choose allowed tools for each MCP server before saving", + ); + expect(mocks.update).not.toHaveBeenCalled(); + + await result.current.saveCustomer({ ...values, mcp_tool_permissions: { [server.server_id]: ["read"] } }); + + expect(mocks.update).toHaveBeenCalledOnce(); + }); + + it.each([["group-a"], [{ name: "group-a" }]])( + "preserves the last toolset ceiling on a retained group server with membership %j", + async (membership) => { + const stored: EndUser = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [], + mcp_access_groups: ["group-a"], + mcp_toolsets: [toolset.toolset_id], + mcp_tool_permissions: {}, + }, + }; + mocks.servers.mockReturnValue({ data: [{ ...server, mcp_access_groups: [membership] }], isSuccess: true }); + mocks.toolsets.mockReturnValue({ data: [toolset], isSuccess: true }); + mocks.accessGroups.mockReturnValue({ data: ["group-a"], isSuccess: true }); + const { result } = renderHook(() => useSaveCustomer(stored)); + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [], accessGroups: ["group-a"], toolsets: [] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow("Choose allowed tools"); + expect(mocks.update).not.toHaveBeenCalled(); + await result.current.saveCustomer({ ...values, mcp_tool_permissions: { [server.server_id]: ["read"] } }); + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ + user_id: stored.user_id, + object_permission: { + ...clearedGrant, + mcp_access_groups: ["group-a"], + mcp_tool_permissions: { [server.server_id]: ["read"] }, + }, + }); + }, + ); + + it.each([ + { + scenario: "removing only the group", + storedGroups: ["group-a"], + storedToolsets: [toolset.toolset_id], + servers: [], + groups: [], + toolsets: [toolset.toolset_id], + }, + { + scenario: "granting a new full-access group", + storedGroups: [], + storedToolsets: [], + servers: [], + groups: ["group-a"], + toolsets: [], + }, + { + scenario: "denying every server", + storedGroups: ["group-a"], + storedToolsets: [toolset.toolset_id], + servers: ["no-mcp-servers"], + groups: ["group-a"], + toolsets: [], + }, + { + scenario: "clearing all grants", + storedGroups: ["group-a"], + storedToolsets: [toolset.toolset_id], + servers: [], + groups: [], + toolsets: [], + }, + { + scenario: "removing an unknown empty group", + storedGroups: ["unknown-group"], + storedToolsets: [toolset.toolset_id], + servers: [], + groups: [], + toolsets: [toolset.toolset_id], + }, + ])( + "allows $scenario without adding a tool map", + async ({ storedGroups, storedToolsets, servers, groups, toolsets }) => { + const stored: EndUser = { + ...customer, + object_permission: { + ...permission, + mcp_servers: [], + mcp_access_groups: storedGroups, + mcp_toolsets: storedToolsets, + mcp_tool_permissions: {}, + }, + }; + mocks.servers.mockReturnValue({ data: [{ ...server, mcp_access_groups: ["group-a"] }], isSuccess: true }); + mocks.toolsets.mockReturnValue({ data: [toolset], isSuccess: true }); + mocks.accessGroups.mockReturnValue({ data: ["group-a"], isSuccess: true }); + const { result } = renderHook(() => useSaveCustomer(stored)); + await result.current.saveCustomer({ + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers, accessGroups: groups, toolsets }, + }); + expect(mocks.update).toHaveBeenCalledExactlyOnceWith({ + user_id: stored.user_id, + object_permission: { + mcp_servers: servers, + mcp_access_groups: groups, + mcp_toolsets: toolsets, + mcp_tool_permissions: {}, + }, + }); + }, + ); + + it("refuses a permission edit while a stored access group's catalog has failed, then saves after it loads", async () => { + const added: MCPServer = { ...server, server_id: "added-server" }; + const stored: EndUser = { ...customer, object_permission: { ...permission, mcp_access_groups: ["group-a"] } }; + mocks.servers.mockReturnValue({ data: [server, added], isSuccess: true }); + mocks.accessGroups.mockReturnValue({ data: undefined, isSuccess: false }); + const { result, rerender } = renderHook(() => useSaveCustomer(stored)); + const values = { + ...toCustomerFormValues(stored), + mcp_servers_and_groups: { servers: [server.server_id, added.server_id], accessGroups: ["group-a"], toolsets: [] }, + mcp_tool_permissions: { [server.server_id]: ["read"], [added.server_id]: ["read"] }, + }; + + await expect(result.current.saveCustomer(values)).rejects.toThrow("MCP catalogs must finish loading"); + expect(mocks.update).not.toHaveBeenCalled(); + mocks.accessGroups.mockReturnValue({ data: ["group-a"], isSuccess: true }); + rerender(); + + await result.current.saveCustomer(values); + + expect(mocks.update).toHaveBeenCalledOnce(); + }); + + it.each([ + { case: "an unrestricted stored server", stored: ["added-server"], toolsetServer: "server-id" }, + { case: "a server its selected toolset covers", stored: [], toolsetServer: "added-server" }, + ])("saves $case without an allowlist", async ({ stored, toolsetServer }) => { + const added: MCPServer = { ...server, server_id: "added-server" }; + const addedToolset: MCPToolset = { ...toolset, tools: [{ server_id: toolsetServer, tool_name: "read" }] }; + const existing: EndUser = { + ...customer, + object_permission: { ...permission, mcp_servers: stored, mcp_tool_permissions: {} }, + }; + mocks.servers.mockReturnValue({ data: [server, added], isSuccess: true }); + mocks.toolsets.mockReturnValue({ data: [addedToolset], isSuccess: true }); + const { result } = renderHook(() => useSaveCustomer(existing)); + + await result.current.saveCustomer({ + ...toCustomerFormValues(existing), + mcp_servers_and_groups: { servers: [added.server_id], accessGroups: [], toolsets: [toolset.toolset_id] }, + }); + + expect(mocks.update).toHaveBeenCalledExactlyOnceWith( + expect.objectContaining({ object_permission: expect.objectContaining({ mcp_toolsets: [toolset.toolset_id] }) }), + ); + }); +}); 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..5230d9f76e0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useSaveCustomer.ts @@ -0,0 +1,50 @@ +import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; +import { useCreateCustomer, useUpdateCustomer } from "@/app/(dashboard)/hooks/customers/useCustomerMutations"; +import { useMCPAccessGroups } from "@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { + buildCustomerPayload, + hasCustomerEditConflict, + serversLosingToolLimits, + type CustomerFormValues, +} from "./customerPayload"; + +export const useSaveCustomer = (existing: EndUser | null) => { + const { data: servers = [], isSuccess: serversReady } = useMCPServers(); + const { data: toolsets = [], isSuccess: toolsetsReady } = useMCPToolsets(); + const { isSuccess: accessGroupsReady } = useMCPAccessGroups(); + const create = useCreateCustomer(); + const update = useUpdateCustomer(); + const saveCustomer = (values: CustomerFormValues, baseline: EndUser | null = existing) => { + const payload = buildCustomerPayload(values, baseline, servers, toolsets); + if (hasCustomerEditConflict(payload, baseline, existing)) { + return Promise.reject( + new Error("Customer permissions or budget changed while editing. Cancel and reopen Settings before saving"), + ); + } + if (!payload.object_permission) return existing ? update.mutateAsync(payload) : create.mutateAsync(payload); + const selection = values.mcp_servers_and_groups; + const stored = baseline?.object_permission; + const catalogs = [ + { needed: true, ready: serversReady }, + { needed: selection.toolsets.length > 0 || (stored?.mcp_toolsets?.length ?? 0) > 0, ready: toolsetsReady }, + { + needed: selection.accessGroups.length > 0 || (stored?.mcp_access_groups?.length ?? 0) > 0, + ready: accessGroupsReady, + }, + ]; + if (catalogs.some(({ needed, ready }) => needed && !ready)) { + return Promise.reject(new Error("MCP catalogs must finish loading before permission changes can be saved")); + } + if (serversLosingToolLimits(values, baseline, servers, toolsets).length > 0) { + return Promise.reject( + new Error( + "Choose allowed tools for each MCP server before saving, or it gets every tool including delete tools", + ), + ); + } + 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..64d2d0ffe44 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.test.ts @@ -0,0 +1,77 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { apiClient } from "@/components/networking"; +import { createCustomer, deleteCustomers, setCustomersBlockedState, 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([ + { + 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 new file mode 100644 index 00000000000..6530df36166 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/customerApi.ts @@ -0,0 +1,38 @@ +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 CustomersBlockedStateInput = { + userIds: 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 setCustomersBlockedState = ( + accessToken: string, + { userIds, blocked }: CustomersBlockedStateInput, +): Promise => + apiClient.post( + blocked ? "/customer/block" : "/customer/unblock", + { + accessToken, + body: { user_ids: userIds }, + }, + ); 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..382ebaed942 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomerMutations.test.tsx @@ -0,0 +1,198 @@ +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, + useSetCustomersBlockedState, + 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.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 new file mode 100644 index 00000000000..933eec760b1 --- /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, setCustomersBlockedState, 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 useSetCustomersBlockedState = () => useCustomerMutation(setCustomersBlockedState); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index e3025266ddb..9e9c0015292 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -42,6 +42,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/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 3a29578c5c2..223ae2416b6 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(); @@ -12,6 +13,7 @@ const mockTeamInfoCall = vi.fn(); const mockUserUpdateUserCall = vi.fn(); const mockFetchMCPServers = vi.fn(); const mockListMCPTools = vi.fn(); +const mockMCPServerCatalog = vi.fn(); const MCP_SERVER = { server_id: "srv-1", server_name: "GitHub MCP", alias: "GitHub MCP" }; @@ -69,7 +71,7 @@ vi.mock("@/components/networking", () => { }); vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ - useMCPServers: () => ({ data: [MCP_SERVER], isLoading: false }), + useMCPServers: () => mockMCPServerCatalog(), })); vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups", () => ({ useMCPAccessGroups: () => ({ data: ["dev-group"], isLoading: false }), @@ -107,6 +109,7 @@ describe("UserInfoView", () => { mockTeamMemberDeleteCall.mockResolvedValue({}); mockUserUpdateUserCall.mockResolvedValue({}); mockFetchMCPServers.mockResolvedValue([MCP_SERVER]); + mockMCPServerCatalog.mockReturnValue({ data: [MCP_SERVER], isLoading: false }); mockListMCPTools.mockResolvedValue({ tools: [{ name: "list_issues", description: "List issues" }] }); }); @@ -330,6 +333,43 @@ describe("UserInfoView", () => { }); describe("MCP permissions", () => { + it.each([ + { field: "user_email", label: "Email", value: "renamed@example.com", catalogLoaded: true }, + { field: "user_alias", label: "User Alias", value: "Renamed user", catalogLoaded: false }, + ])( + "preserves shared-alias grants when only $field changes with catalogLoaded=$catalogLoaded", + async ({ field, label, value, catalogLoaded }) => { + const servers = [ + { server_id: "server-a", server_name: "Server A", alias: "shared" }, + { server_id: "server-b", server_name: "Server B", alias: "shared" }, + ]; + const permission = { + mcp_servers: ["server-a"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { shared: ["list_issues"] }, + }; + mockUserGetInfoV2.mockResolvedValue({ ...MOCK_USER_DATA, object_permission: permission }); + mockMCPServerCatalog.mockReturnValue({ data: catalogLoaded ? servers : [], isLoading: !catalogLoaded }); + if (catalogLoaded) mockFetchMCPServers.mockResolvedValue(servers); + else mockFetchMCPServers.mockReturnValue(new Promise(() => {})); + render(); + + fireEvent.change(await screen.findByRole("textbox", { name: label, exact: true }), { target: { value } }); + await userEvent.click(screen.getByRole("button", { name: "Save Changes", exact: true })); + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalledOnce()); + + const [, payload] = mockUserUpdateUserCall.mock.calls[0]; + expect(payload[field]).toBe(value); + expect(payload).not.toHaveProperty("object_permission"); + await userEvent.click(await screen.findByRole("button", { name: "Edit Settings", exact: true })); + expect(await screen.findByRole("textbox", { name: label, exact: true })).toHaveValue(value); + await userEvent.click(screen.getByRole("button", { name: "Save Changes", exact: true })); + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalledTimes(2)); + expect(mockUserUpdateUserCall.mock.calls[1][1]).not.toHaveProperty("object_permission"); + }, + ); + it("should render the user's MCP entitlements in read mode", async () => { const user = userEvent.setup(); render(); @@ -347,7 +387,7 @@ describe("UserInfoView", () => { expect(await screen.findByText("list_issues")).toBeInTheDocument(); }); - it("should nest MCP entitlements under object_permission when an admin saves", async () => { + it("omits unchanged MCP entitlements when an admin saves", async () => { const user = userEvent.setup(); render(); @@ -362,13 +402,7 @@ describe("UserInfoView", () => { expect(token).toBe("test-token"); expect(roleArg).toBeNull(); expect(payload.user_id).toBe("user-123"); - const expectedObjectPermission = { - mcp_servers: ["srv-1"], - mcp_access_groups: ["dev-group"], - mcp_toolsets: [], - mcp_tool_permissions: { "srv-1": ["list_issues"] }, - }; - expect(payload.object_permission).toEqual(expectedObjectPermission); + expect(payload).not.toHaveProperty("object_permission"); expect(payload).not.toHaveProperty("mcp_servers_and_groups"); expect(payload).not.toHaveProperty("mcp_tool_permissions"); expect(payload).not.toHaveProperty("mcp_servers"); @@ -417,9 +451,35 @@ describe("UserInfoView", () => { }); const [, payload] = mockUserUpdateUserCall.mock.calls[0]; - expect(payload.object_permission.mcp_tool_permissions).toEqual({ - "srv-1": ["list_issues"], - "srv-via-group": ["read_only"], + expect(payload).not.toHaveProperty("object_permission"); + }); + + it("retains an independent tool-only server grant when another server's tools are edited", async () => { + const user = userEvent.setup(); + const servers = [MCP_SERVER, { server_id: "server-b", server_name: "Server B" }]; + const permission = { + mcp_servers: ["srv-1"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "srv-1": ["read_direct"], "server-b": ["read_independent"] }, + }; + mockUserGetInfoV2.mockResolvedValue({ ...MOCK_USER_DATA, object_permission: permission }); + mockFetchMCPServers.mockResolvedValue(servers); + mockMCPServerCatalog.mockReturnValue({ data: servers, isLoading: false }); + mockListMCPTools.mockImplementation((_token: string, serverId: string) => + Promise.resolve({ tools: [{ name: serverId === "srv-1" ? "read_direct" : "read_independent" }] }), + ); + render(); + + await user.click((await screen.findAllByRole("radio", { name: "Flat List" }))[0]); + await user.click(await screen.findByRole("checkbox", { name: "read_direct", exact: true })); + await user.click(screen.getByRole("button", { name: "Save Changes", exact: true })); + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalledOnce()); + + const [, payload] = mockUserUpdateUserCall.mock.calls[0]; + expect(payload.object_permission).toEqual({ + ...permission, + mcp_tool_permissions: { "srv-1": [], "server-b": ["read_independent"] }, }); }); @@ -442,6 +502,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 }, @@ -495,6 +561,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. @@ -507,22 +596,30 @@ describe("extractMcpEntitlement", () => { expect(result?.mcp_tool_permissions).toEqual({ "srv-deleted": ["read"] }); }); - it.each([ - ["selected server first", ["srv-shared-a", "srv-shared-b"]], - ["deselected server first", ["srv-shared-b", "srv-shared-a"]], - ])("keeps a shared-name allowlist while any server it names is granted (%s)", (_label, order) => { - // Names are not unique and the gateway unions a name key into EVERY server answering to it, so - // this one entry restricts both. Resolving to the first match would drop it whenever the catalog - // happened to return the deselected server first, stripping the restriction from the one still - // granted, which is a widening that reproduces on one deployment and not another. - const catalog = [ - { server_id: "srv-shared-a", server_name: "shared", alias: null }, - { server_id: "srv-shared-b", server_name: "shared", alias: null }, - ] as any; - const ordered = order.map((id) => catalog.find((server: any) => server.server_id === id)); + it.each([false, true])("limits a shared alias to retained servers regardless of catalog order (%s)", (reverse) => { + const catalog: MCPServer[] = [ + { ...auditFields, server_id: "srv-shared-a", server_name: "kept", alias: "shared" }, + { ...auditFields, server_id: "srv-shared-b", server_name: "removed", alias: "shared" }, + ]; + const ordered = reverse ? catalog.toReversed() : catalog; + const permissions = { + "srv-shared-a": ["read"], + kept: ["write"], + shared: ["list"], + "srv-shared-b": ["delete"], + removed: ["admin"], + }; + const result = extractMcpEntitlement(form({ servers: ["srv-shared-a"] }, permissions), ordered); + expect(result?.mcp_tool_permissions).toEqual({ "srv-shared-a": ["read", "write", "list"] }); + }); - const result = extractMcpEntitlement(form({ servers: ["srv-shared-a"] }, { shared: ["read"] }), ordered as any); - expect(result?.mcp_tool_permissions).toEqual({ shared: ["read"] }); + it("keeps shared-alias restrictions on servers retained through access groups", () => { + const catalog: MCPServer[] = [ + { ...auditFields, server_id: "kept", alias: "shared", mcp_access_groups: ["retained"] }, + { ...auditFields, server_id: "removed", alias: "shared" }, + ]; + const result = extractMcpEntitlement(form({ accessGroups: ["retained"] }, { shared: ["*"] }), catalog); + expect(result?.mcp_tool_permissions).toEqual({ kept: ["*"] }); }); it("drops a shared-name allowlist once no server it names is granted", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index 2056142b50a..2839cae82cd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -42,7 +42,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import MCPServerPermissions from "@/components/permissions/MCPServerPermissions"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; -import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; +import { extractMcpEntitlementUpdate } from "@/components/mcp_server_management/mcpEntitlement"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab"; @@ -321,7 +321,12 @@ export default function UserInfoView({ try { if (!accessToken || !userData) return; - const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers, allMcpToolsets); + const mcpEntitlement = extractMcpEntitlementUpdate( + formValues, + userData.object_permission, + allMcpServers, + allMcpToolsets, + ); const userFields = Object.fromEntries( Object.entries(formValues).filter( ([field]) => field !== "mcp_servers_and_groups" && field !== "mcp_tool_permissions", diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index b811ab8903c..a30ff73e37f 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, @@ -350,6 +351,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 6e2c48afaab..71f0c248952 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( @@ -244,3 +258,42 @@ describe("MCPServerSelector scoping", () => { expect(optionByLabel("Wiki A")).toBeDefined(); }); }); + +describe("MCPServerSelector name collisions", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseMCPServers.mockReturnValue({ + data: [ + { server_id: "shared", server_name: "Shared" }, + { server_id: "srv-x", server_name: "Other" }, + ], + isLoading: false, + } as any); + mockUseMCPAccessGroups.mockReturnValue({ data: ["shared", "alias-a"], isLoading: false } as any); + mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false } as any); + }); + + it.each(["shared", "alias-a"])( + "keeps the stored server %s a server when it matches an access group name", + async (server) => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders( + , + ); + + await openSelector(user); + await user.click(optionByLabel("Other")!); + + expect(onChange).toHaveBeenCalledWith({ servers: [server, "srv-x"], accessGroups: [], toolsets: [] }); + }, + ); + + it("shows a stored server alias by its name", () => { + renderWithProviders( + , + ); + + expect(screen.getByLabelText("alias-a")).toBeInTheDocument(); + }); +}); 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 2d78d757aab..073f995f3f5 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -8,6 +8,7 @@ import type { MCPServer } from "@/components/mcp_tools/types"; import { accessGroupNamesOf } from "./effectiveMcpServers"; interface MCPServerSelectorProps { + id?: string; onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; value?: { servers: string[]; @@ -25,6 +26,14 @@ interface MCPServerSelectorProps { } const TOOLSET_PREFIX = "toolset:"; +const ACCESS_GROUP_PREFIX = "access-group:"; +const SERVER_PREFIX = "server:"; +const SENTINELS: readonly string[] = [ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL]; + +const serverValue = (server: string): string => (SENTINELS.includes(server) ? server : `${SERVER_PREFIX}${server}`); + +const valuesWithPrefix = (values: readonly string[], prefix: string): string[] => + values.filter((value) => value.startsWith(prefix)).map((value) => value.slice(prefix.length)); export const serversWithinAccessGroups = ( servers: readonly MCPServer[], @@ -41,6 +50,7 @@ export const serversWithinAccessGroups = ( }; const MCPServerSelector: React.FC = ({ + id, onChange, value, className, @@ -58,19 +68,18 @@ const MCPServerSelector: React.FC = ({ const loading = serversLoading || groupsLoading || toolsetsLoading; - const accessGroupSet = new Set(accessGroups); const visibleServers = serversWithinAccessGroups(mcpServers, value?.accessGroups ?? [], value?.servers ?? []); // Combine options: access groups + servers + toolsets const options = [ ...accessGroups.map((group) => ({ label: group, - value: group, + value: `${ACCESS_GROUP_PREFIX}${group}`, description: "Access Group", })), ...visibleServers.map((server) => ({ label: `${server.server_name || server.server_id} (${server.server_id})`, - value: server.server_id, + value: serverValue(server.server_id), description: "MCP Server", })), ...toolsets.map((toolset) => ({ @@ -80,13 +89,24 @@ const MCPServerSelector: React.FC = ({ })), ]; - // Flatten value for Select — prefix toolset IDs const selectedValues = [ - ...(value?.servers || []), - ...(value?.accessGroups || []), + ...(value?.servers || []).map(serverValue), + ...(value?.accessGroups || []).map((group) => `${ACCESS_GROUP_PREFIX}${group}`), ...(value?.toolsets || []).map((id) => `${TOOLSET_PREFIX}${id}`), ]; + const listedValues = new Set(options.map((option) => option.value)); + const unlistedSelections = [ + ...(value?.servers || []) + .filter((server) => !SENTINELS.includes(server)) + .map((server) => ({ label: server, value: serverValue(server), description: "MCP Server" })), + ...(value?.accessGroups || []).map((group) => ({ + label: group, + value: `${ACCESS_GROUP_PREFIX}${group}`, + description: "Access Group", + })), + ].filter((option) => !listedValues.has(option.value)); + const hasNoMcpServersSelected = allowNoMcpServers && selectedValues.includes(NO_MCP_SERVERS_SENTINEL); const hasAllProxyMcpServersSelected = selectedValues.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); @@ -101,13 +121,11 @@ const MCPServerSelector: React.FC = ({ onChange({ servers: [NO_MCP_SERVERS_SENTINEL], accessGroups: [], toolsets: [] }); return; } - const toolsetsSelected = selected - .filter((v) => v.startsWith(TOOLSET_PREFIX)) - .map((v) => v.slice(TOOLSET_PREFIX.length)); - const rest = selected.filter((v) => !v.startsWith(TOOLSET_PREFIX)); - const servers = rest.filter((v) => !accessGroupSet.has(v)); - const accessGroupsSelected = rest.filter((v) => accessGroupSet.has(v)); - onChange({ servers, accessGroups: accessGroupsSelected, toolsets: toolsetsSelected }); + onChange({ + servers: [...selected.filter((v) => SENTINELS.includes(v)), ...valuesWithPrefix(selected, SERVER_PREFIX)], + accessGroups: valuesWithPrefix(selected, ACCESS_GROUP_PREFIX), + toolsets: valuesWithPrefix(selected, TOOLSET_PREFIX), + }); }; const selectOptions: MultiSelectOption[] = [ @@ -117,7 +135,7 @@ const MCPServerSelector: React.FC = ({ ...(allowNoMcpServers ? [{ label: "No MCP Servers", value: NO_MCP_SERVERS_SENTINEL, description: "Block all" }] : []), - ...options.map((opt) => ({ + ...[...options, ...unlistedSelections].map((opt) => ({ ...opt, disabled: hasNoMcpServersSelected || hasAllProxyMcpServersSelected, })), @@ -126,6 +144,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"] }, "edit"); + }); + + 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, + preserveUnrestrictedServers: [mockServerId], + }; + const { rerender } = renderWithProviders(); + + expect(await screen.findByText("delete_record")).toBeInTheDocument(); + expect(onChange).not.toHaveBeenCalled(); + + rerender(); + + await waitFor(() => { + expect(onChange).toHaveBeenCalledExactlyOnceWith({ [newServerId]: ["list_records"] }, "default"); + }); + }); + 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/mcpEntitlement.ts b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts index 66f59f5b71d..af69928918d 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,13 @@ -import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import type { ObjectPermission } from "@/components/object_permission_types"; +import { + accessGroupNamesOf, + mcpAllowedToolsFor, + mcpServersForIdentifier, + mcpToolPermissionKeysFor, + resolveEffectiveMcpServers, +} from "./effectiveMcpServers"; export interface McpEntitlementUpdate { mcp_servers: string[]; @@ -18,45 +26,31 @@ const asToolPermissions = (value: unknown): Record => { ); }; -const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): boolean => - server.server_id === identifier || server.server_name === identifier || server.alias === identifier; +const formMcpGrant = (formValues: Record): McpEntitlementUpdate | null => { + const selection = formValues.mcp_servers_and_groups; + if (selection === null || typeof selection !== "object") return null; + const { servers, accessGroups, toolsets } = selection as Record; + return { + mcp_servers: asStringArray(servers), + mcp_access_groups: asStringArray(accessGroups), + mcp_toolsets: asStringArray(toolsets), + mcp_tool_permissions: asToolPermissions(formValues.mcp_tool_permissions), + }; +}; -/** - * The `object_permission` a save sends, derived from what the editor currently shows. - * - * A tool allowlist is what narrows a grant and an absent one reads as no restriction, so dropping - * an entry is the direction that widens. An entry is kept while its own server is still reachable, - * directly or through a retained access group or toolset, and dropped once nothing reaches it, which - * is what makes removing a grant actually remove it. Reachability is resolved per server rather than - * per selection: the gateway treats every allowlist key as an independent server grant, so keeping - * every key because some unrelated group survived would leave a deselected server callable. - * - * The catalog carries `mcp_access_groups` on each server and `tools[].server_id` on each toolset, - * 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. - * `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 - * on catalog order and could drop a restriction that was also covering a server still granted. A key - * that resolves to nothing is kept too, since a server we cannot identify is one we cannot confirm - * was deselected; that also covers a catalog that has not loaded or failed to load, where every key - * is unresolvable and nothing is pruned. - */ export const extractMcpEntitlement = ( formValues: Record, allServers: MCPServer[], allToolsets: MCPToolset[] = [], ): McpEntitlementUpdate | null => { - const selection = formValues.mcp_servers_and_groups; - if (selection === null || typeof selection !== "object") return null; - - const { servers, accessGroups, toolsets } = selection as Record; - const mcpServers = asStringArray(servers); - const mcpAccessGroups = asStringArray(accessGroups); - const mcpToolsets = asStringArray(toolsets); + const formGrant = formMcpGrant(formValues); + if (!formGrant) return null; + const { + mcp_servers: mcpServers, + mcp_access_groups: mcpAccessGroups, + mcp_toolsets: mcpToolsets, + mcp_tool_permissions: toolPermissions, + } = formGrant; const grantsEveryServer = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL) || mcpToolsets.some((toolsetId) => !allToolsets.some((toolset) => toolset.toolset_id === toolsetId)); @@ -66,26 +60,140 @@ 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)) || - (server.mcp_access_groups ?? []).some((group) => mcpAccessGroups.includes(group)) || + directServerIds.has(server.server_id) || + accessGroupNamesOf(server).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); }; + const splitKeys = new Set( + Object.keys(toolPermissions).filter((key) => { + const named = mcpServersForIdentifier(allServers, key); + return !grantsEveryServer && named.length > 1 && named.some((server) => !grants(server)); + }), + ); + const retainedServerIds = new Set( + Array.from(splitKeys).flatMap((key) => + mcpServersForIdentifier(allServers, key) + .filter(grants) + .map((server) => server.server_id), + ), + ); + const keepPermissionKey = (key: string): boolean => { + if (splitKeys.has(key)) return false; + const named = mcpServersForIdentifier(allServers, key); + if (named.length === 1 && retainedServerIds.has(named[0].server_id)) return false; + return grantsEveryServer || grantsServerNamedBy(key); + }; + const retainedPermissions = Object.fromEntries( + Object.entries(toolPermissions).filter(([key]) => keepPermissionKey(key)), + ); + const canonicalPermissions = Object.fromEntries( + allServers + .filter((server) => retainedServerIds.has(server.server_id)) + .map((server) => [server.server_id, [...(mcpAllowedToolsFor(server, toolPermissions, allServers) ?? [])]]), + ); + return { mcp_servers: mcpServers, mcp_access_groups: mcpAccessGroups, mcp_toolsets: mcpToolsets, - mcp_tool_permissions: Object.fromEntries( - Object.entries(asToolPermissions(formValues.mcp_tool_permissions)).filter( - ([permissionKey]) => grantsEveryServer || grantsServerNamedBy(permissionKey), - ), - ), + mcp_tool_permissions: { ...retainedPermissions, ...canonicalPermissions }, }; }; + +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 independentToolPermissions = ( + storedGrant: McpEntitlementUpdate, + toolPermissions: Record, + allServers: MCPServer[], + allToolsets: MCPToolset[], +): Record => { + if (storedGrant.mcp_servers.includes(NO_MCP_SERVERS_SENTINEL)) return {}; + const storedSelection = { + allServers, + selectedServers: storedGrant.mcp_servers, + selectedAccessGroups: storedGrant.mcp_access_groups, + selectedToolsets: storedGrant.mcp_toolsets, + toolsets: allToolsets, + toolPermissions: {}, + }; + const selectedServers = resolveEffectiveMcpServers(storedSelection); + const selectedServerIds = new Set(selectedServers.map(({ server }) => server.server_id)); + const keyOnlyServers = Object.keys(storedGrant.mcp_tool_permissions) + .flatMap((key) => mcpServersForIdentifier(allServers, key)) + .filter((server) => !selectedServerIds.has(server.server_id)); + const keyOnlyServerIds = new Set(keyOnlyServers.map((server) => server.server_id)); + const namesOnlyKeyOnlyServers = (key: string): boolean => { + const named = mcpServersForIdentifier(allServers, key); + return named.length > 0 && named.every((server) => keyOnlyServerIds.has(server.server_id)); + }; + const independentPermissions = Object.fromEntries( + Object.entries(toolPermissions).filter(([key]) => namesOnlyKeyOnlyServers(key)), + ); + const canonicalPermissions = Object.fromEntries( + keyOnlyServers + .filter((server) => mcpToolPermissionKeysFor(server, independentPermissions, allServers).length === 0) + .flatMap((server) => { + const allowed = mcpAllowedToolsFor(server, toolPermissions, allServers); + return allowed ? [[server.server_id, [...allowed]]] : []; + }), + ); + return { ...independentPermissions, ...canonicalPermissions }; +}; + +export const extractMcpEntitlementUpdate = ( + formValues: Record, + existingPermission: ObjectPermission | null | undefined, + allServers: MCPServer[], + allToolsets: MCPToolset[] = [], +): McpEntitlementUpdate | null => { + const formGrant = formMcpGrant(formValues); + if (!formGrant) return null; + const storedGrant = { + mcp_servers: existingPermission?.mcp_servers ?? [], + mcp_access_groups: existingPermission?.mcp_access_groups ?? [], + mcp_toolsets: existingPermission?.mcp_toolsets ?? [], + mcp_tool_permissions: existingPermission?.mcp_tool_permissions ?? {}, + }; + const storedIdentity = grantIdentity(storedGrant); + if (grantIdentity(formGrant) === storedIdentity) return null; + const extracted = extractMcpEntitlement(formValues, allServers, allToolsets); + if (!extracted) return null; + const grant = { + ...extracted, + mcp_tool_permissions: formGrant.mcp_servers.includes(NO_MCP_SERVERS_SENTINEL) + ? {} + : { + ...extracted.mcp_tool_permissions, + ...independentToolPermissions(storedGrant, formGrant.mcp_tool_permissions, allServers, allToolsets), + }, + }; + return grantIdentity(grant) === storedIdentity ? null : grant; +}; diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 18f136c3b28..fc9b40f3998 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, model and MCP access", teams: "Create and manage teams for access control", organizations: "Manage organizations and their members", projects: "Manage projects within teams", diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx index e34359687d9..faafc4c23c4 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.test.tsx @@ -157,6 +157,39 @@ describe("MCPServerPermissions", () => { expect(screen.queryByText(/tool/)).not.toBeInTheDocument(); }); + it("keeps the existing key display for empty tool allowlists", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 }, + ]); + render( + , + ); + expect(await screen.findByText(/DW_MCP/)).toBeInTheDocument(); + expect(screen.queryByText("No tools allowed")).not.toBeInTheDocument(); + }); + + it("labels an empty tool allowlist as allowing no tools", async () => { + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { server_id: mockServerId1, server_name: mockServerName1, alias: mockServerName1 }, + ]); + + render( + , + ); + + expect(await screen.findByText(/DW_MCP/)).toBeInTheDocument(); + expect(screen.getByText("No tools allowed")).toBeInTheDocument(); + }); + it("should display access groups correctly", async () => { /** * Tests that access groups are displayed with the correct styling @@ -464,6 +497,46 @@ describe("MCPServerPermissions", () => { />, ); await waitFor(() => expect(screen.getByText("Blocked")).toHaveAttribute("data-variant", "destructive")); + expect( + screen.getByText("No MCP servers. This key is blocked from all MCP servers, including its team's servers"), + ).toBeInTheDocument(); + }); + + it.each(["server", "tool permission"])("shows no customer access for an unsupported %s sentinel", (source) => { + render( + , + ); + + expect(screen.getByText("No MCP access")).toBeInTheDocument(); + expect(screen.getByText("0")).toBeInTheDocument(); + expect(screen.queryByText("All")).not.toBeInTheDocument(); + }); + + it.each(["direct", "tool permission"])( + "shows the customer's concrete %s grant beside an unsupported sentinel", + (source) => { + render( + , + ); + + expect(screen.getByText(mockServerId1)).toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.queryByText("All")).not.toBeInTheDocument(); + expect(screen.queryByText("No MCP access")).not.toBeInTheDocument(); + }, + ); + + it("uses customer-specific blocked copy without implying a team relationship", () => { + render(); + expect(screen.getByText("No MCP servers. This customer is blocked from all MCP servers")).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..3b8039aff38 100644 --- a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx +++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx @@ -10,6 +10,7 @@ import { InheritedGrant, inheritedGrantTooltip } from "./inheritedGrants"; interface MCPServerPermissionsProps { mcpServers: string[]; + entityType?: "key" | "customer"; mcpAccessGroups?: string[]; mcpToolPermissions?: Record; mcpToolsets?: string[]; @@ -19,6 +20,7 @@ interface MCPServerPermissionsProps { export function MCPServerPermissions({ mcpServers, + entityType = "key", mcpAccessGroups = [], mcpToolPermissions = {}, mcpToolsets = [], @@ -54,10 +56,12 @@ export function MCPServerPermissions({ }); }; - const directServerIds = mcpServers.filter( + const serverIdentifiers = + entityType === "customer" ? [...new Set([...mcpServers, ...Object.keys(mcpToolPermissions)])] : mcpServers; + const directServerIds = serverIdentifiers.filter( (server) => server !== NO_MCP_SERVERS_SENTINEL && server !== ALL_PROXY_MCP_SERVERS_SENTINEL, ); - const inheritedOnlyServers = inheritedMcpServers.filter((grant) => !mcpServers.includes(grant.id)); + const inheritedOnlyServers = inheritedMcpServers.filter((grant) => !directServerIds.includes(grant.id)); const serverIdCount = directServerIds.length + inheritedOnlyServers.length; // Fetch MCP server details when component mounts @@ -118,7 +122,8 @@ export function MCPServerPermissions({ }; const blocksAllMcpServers = mcpServers.includes(NO_MCP_SERVERS_SENTINEL); - const grantsAllProxyMcpServers = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); + const hasAllProxySentinel = serverIdentifiers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL); + const grantsAllProxyMcpServers = entityType === "key" && hasAllProxySentinel; const mergedItems = [ ...directServerIds.map((server) => ({ type: "server", value: server, tooltip: `Full ID: ${server}` })), @@ -145,7 +150,9 @@ export function MCPServerPermissions({

- No MCP servers — this key is blocked from all MCP servers, including its team's servers + {entityType === "key" + ? "No MCP servers. This key is blocked from all MCP servers, including its team's servers" + : "No MCP servers. This customer is blocked from all MCP servers"}

) : grantsAllProxyMcpServers ? ( @@ -158,6 +165,7 @@ export function MCPServerPermissions({ {mergedItems.map((item, index) => { const toolsForServer = item.type === "server" ? getToolPermissionsFor(item.value) : undefined; const hasToolRestrictions = toolsForServer && toolsForServer.length > 0; + const allowsNoTools = entityType === "customer" && toolsForServer?.length === 0; const isExpanded = expandedServers.has(item.value); return ( @@ -190,6 +198,7 @@ export function MCPServerPermissions({ )}
+ {allowsNoTools && No tools allowed} {hasToolRestrictions && (
{toolsForServer.length} @@ -285,7 +294,11 @@ export function MCPServerPermissions({ ) : (
-

No MCP servers, access groups, or toolsets configured

+

+ {entityType === "customer" && hasAllProxySentinel + ? "No MCP access" + : "No MCP servers, access groups, or toolsets configured"} +

)}
diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx index 3a618114613..29d2eebec9a 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.test.tsx @@ -62,6 +62,17 @@ describe("ModelsCell", () => { expect(await screen.findByText("d")).toBeInTheDocument(); }); + it.each([[], ["all-proxy-models"], ["*"]])("uses the customer model scope label for %j", (models) => { + render(); + expect(screen.getByText("All models the key can use")).toBeInTheDocument(); + }); + + it("keeps the literal wildcard label on existing pages", () => { + render(); + expect(screen.getByText("*")).toBeInTheDocument(); + expect(screen.queryByText("All Proxy Models")).not.toBeInTheDocument(); + }); + it("labels the all-proxy-models wildcard", () => { render(); expect(screen.getByText("All Proxy Models")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx index 81928d6fe87..30aad922e6a 100644 --- a/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx +++ b/ui/litellm-dashboard/src/components/shared/table_cells/models_cell.tsx @@ -9,21 +9,22 @@ import { CellTooltip } from "./cell_tooltip"; interface ModelsCellProps { models: string[] | null | undefined; maxVisible?: number; + unrestrictedLabel?: string; allowedRoutes?: string[] | null; keyType?: string | null; } const WILDCARD_MODEL = "all-proxy-models"; -const formatModel = (model: string): string => { - if (model === WILDCARD_MODEL) { - return "All Proxy Models"; +const formatModel = (model: string, unrestrictedLabel?: string): string => { + if (model === WILDCARD_MODEL || (model === "*" && unrestrictedLabel !== undefined)) { + return unrestrictedLabel ?? "All Proxy Models"; } const name = getModelDisplayName(model); return name.length > 30 ? `${name.slice(0, 30)}...` : name; }; -export function ModelsCell({ models, maxVisible = 3, allowedRoutes, keyType }: ModelsCellProps) { +export function ModelsCell({ models, maxVisible = 3, unrestrictedLabel, allowedRoutes, keyType }: ModelsCellProps) { if (!Array.isArray(models) || models.length === 0) { const scope = deriveKeyModelScope(allowedRoutes, keyType); if (!scope.hasModelAccess) { @@ -38,7 +39,7 @@ export function ModelsCell({ models, maxVisible = 3, allowedRoutes, keyType }: M /> ); } - return All Proxy Models; + return {unrestrictedLabel ?? "All Proxy Models"}; } const visible = models.slice(0, maxVisible); @@ -48,7 +49,7 @@ export function ModelsCell({ models, maxVisible = 3, allowedRoutes, keyType }: M
{visible.map((model, index) => ( - {formatModel(model)} + {formatModel(model, unrestrictedLabel)} ))} {overflow.length > 0 && ( @@ -56,7 +57,7 @@ export function ModelsCell({ models, maxVisible = 3, allowedRoutes, keyType }: M content={
{overflow.map((model, index) => ( - {formatModel(model)} + {formatModel(model, unrestrictedLabel)} ))}
} diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index c82a4823dfd..b271737fe07 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -1,6 +1,11 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { toast } from "@/lib/toast"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import type { KeyResponse } from "../key_team_helpers/key_list"; +import type { MCPServer } from "../mcp_tools/types"; +import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { toKeyEditFormValues, toSubmittedValues } from "./keyEditFormValues"; vi.mock("@/lib/toast", () => ({ toast: { error: vi.fn(), success: vi.fn() }, @@ -228,7 +233,7 @@ const baseKeyData = { next_rotation_at: null as any, }; -const renderView = (premiumUser: boolean) => { +const renderView = (premiumUser: boolean, keyData: KeyResponse = baseKeyData as KeyResponse) => { // Configure the mock for this test mockUseAuthorized.mockReturnValue({ accessToken: "access_abc", @@ -242,18 +247,13 @@ const renderView = (premiumUser: boolean) => { }); return render( - {}} - keyData={baseKeyData as any} - onKeyDataUpdate={() => {}} - teams={[]} - />, + {}} keyData={keyData} onKeyDataUpdate={() => {}} teams={[]} />, ); }; beforeEach(() => { vi.clearAllMocks(); + vi.mocked(useMCPServers).mockReturnValue({ data: [] } as ReturnType); (globalThis as any).__TEST_FORM_VALUES = undefined; }); @@ -375,6 +375,85 @@ describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { }); }); +describe("KeyInfoView handleKeyUpdate shared MCP aliases", () => { + it.each([ + { + edit: "only renaming the key", + storedServers: ["server-a"], + storedToolPermissions: { shared: ["read"] }, + selectedServers: ["server-a"], + expectedToolPermissions: { shared: ["read"] }, + }, + { + edit: "removing an explicitly selected server", + storedServers: ["server-a", "server-b"], + storedToolPermissions: { shared: ["read"] }, + selectedServers: ["server-a"], + expectedToolPermissions: { "server-a": ["read"] }, + }, + { + edit: "choosing No MCP Servers", + storedServers: ["server-a"], + storedToolPermissions: { "server-b": ["read"] }, + selectedServers: [NO_MCP_SERVERS_SENTINEL], + expectedToolPermissions: {}, + }, + ])( + "preserves unrelated permissions when $edit", + async ({ storedServers, storedToolPermissions, selectedServers, expectedToolPermissions }) => { + const servers: MCPServer[] = ["server-a", "server-b"].map((server_id) => ({ + server_id, + alias: "shared", + created_at: "", + created_by: "", + updated_at: "", + updated_by: "", + })); + vi.mocked(useMCPServers).mockReturnValue({ data: servers } as ReturnType); + const permission = { + mcp_servers: storedServers, + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: storedToolPermissions, + vector_stores: ["vector-store"], + agents: ["agent"], + agent_access_groups: ["agent-group"], + skills: ["skill"], + search_tools: ["search-tool"], + }; + const keyData = { ...baseKeyData, object_permission: permission } as KeyResponse; + renderView(true, keyData); + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as { __TEST_FORM_VALUES?: Record }).__TEST_FORM_VALUES = toSubmittedValues( + { + ...toKeyEditFormValues(keyData), + key_alias: "Renamed key", + mcp_servers_and_groups: { servers: selectedServers, accessGroups: [], toolsets: [] }, + }, + { canViewPolicies: true, canViewPrompts: true }, + ); + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => + expect(keyUpdateCallMock).toHaveBeenCalledExactlyOnceWith( + "access_abc", + expect.objectContaining({ + key: keyData.token, + key_alias: "Renamed key", + object_permission: { + ...permission, + mcp_servers: selectedServers, + mcp_tool_permissions: expectedToolPermissions, + }, + }), + ), + ); + }, + ); +}); + describe("KeyInfoView handleKeyUpdate cache sync", () => { it("should invalidate every cached key query so the list and detail views re-read the saved key", async () => { keyUpdateCallMock.mockResolvedValueOnce({ diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 7b924243fb5..148eedffd04 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -45,7 +45,7 @@ import { useQueryClient } from "@tanstack/react-query"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; -import { extractMcpEntitlement } from "../mcp_server_management/mcpEntitlement"; +import { extractMcpEntitlementUpdate } from "../mcp_server_management/mcpEntitlement"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; @@ -269,7 +269,12 @@ export default function KeyInfoView({ delete formValues.vector_stores; } - const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers ?? [], allMcpToolsets ?? []); + const mcpEntitlement = extractMcpEntitlementUpdate( + formValues, + currentKeyData.object_permission, + allMcpServers ?? [], + allMcpToolsets ?? [], + ); if (mcpEntitlement) { // Without a catalog the grants an allowlist key still has are unresolvable, so nothing is // pruned and a revocation would save as a no-op while reporting success. Refuse instead. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 733a093bcac..f50842e3384 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4174,6 +4174,47 @@ 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 with the same fields as `/customer/update`. Only proxy admins can call it. + * + * Each item runs in its own locked transaction, so a failed item rolls back only its own writes and + * the rest still apply. Results come back in request order with `user_id`, `success`, and either + * `updated_customer` or `error`. + * + * Omitted or null permission fields keep their stored values, supplied lists replace that field, + * and empty lists or objects clear it. + * + * `max_budget` only updates a budget used by that customer alone. Shared and default budgets must be + * edited through the budget endpoints. Sending both `budget_id` and `max_budget` in any item, even + * as null, fails the whole request with HTTP 422. + * + * 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; @@ -4503,12 +4544,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: * ``` @@ -4852,6 +4894,47 @@ 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 with the same fields as `/customer/update`. Only proxy admins can call it. + * + * Each item runs in its own locked transaction, so a failed item rolls back only its own writes and + * the rest still apply. Results come back in request order with `user_id`, `success`, and either + * `updated_customer` or `error`. + * + * Omitted or null permission fields keep their stored values, supplied lists replace that field, + * and empty lists or objects clear it. + * + * `max_budget` only updates a budget used by that customer alone. Shared and default budgets must be + * edited through the budget endpoints. Sending both `budget_id` and `max_budget` in any item, even + * as null, fails the whole request with HTTP 422. + * + * 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; @@ -5181,12 +5264,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: * ``` @@ -28019,6 +28103,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 @@ -31247,6 +31350,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. @@ -57003,6 +57116,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?: { @@ -58024,7 +58170,40 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["BlockUsersResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + 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 */ @@ -58409,7 +58588,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/json": components["schemas"]["UnblockUsersResponse"]; }; }; /** @description Validation Error */