From 5c8b017c89f6c7adde4e7151de9d3412b833b674 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 29 Jul 2026 17:11:48 -0700 Subject: [PATCH] feat(mcp): enforce per-user MCP tool-call entitlements in the auth module The MCP gateway resolved a caller's allowed servers and per-server tool allowlists from the key, the team, the end user and the agent, but never from the internal user row, so an admin had no way to bound what a person may call across every key they hold. Anything the key allowed went through The internal user now carries the same object_permission an admin already attaches to a key or a team, and the resolver applies it as a ceiling: the caller ends up with the intersection of what the key allows and what the user allows, so adding a user entitlement can only narrow, never widen. A level that names no server and no tool places no ceiling, which keeps every existing deployment on its current behavior /user/new and /user/update accept object_permission and reuse the same create-or-update helper the team endpoints use, so the row is written once and the three cached views of it (the user row, the object-permission link and the permission itself) are invalidated on write. Clearing it with an empty object now really unlinks the permission instead of being swallowed as an empty value A row that cannot be read at all places no ceiling, but a row that names a permission the database cannot return denies the call rather than falling through to the wider set, so a partial outage cannot hand out access the admin withheld The users page grows the MCP servers, access groups, toolsets and per-server tool pickers the key and team pages already have. A save keeps a tool allowlist whenever an access group or toolset the admin retained could still supply that server, since an allowlist is what narrows a grant and an absent one reads as no restriction; it drops the allowlist once nothing indirect survives to supply the server, so removing a grant really removes it --- .../mcp_server/auth/user_api_key_auth_mcp.py | 231 +++++++++++- .../mcp_server/mcp_server_manager.py | 5 + litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 3 +- .../proxy/common_utils/user_api_key_cache.py | 22 ++ .../internal_user_endpoints.py | 105 +++++- .../auth/test_user_api_key_auth_mcp.py | 343 ++++++++++++++++++ .../test_internal_user_endpoints.py | 330 +++++++++++++++++ .../users/_components/user_edit_view.tsx | 66 +++- .../view_users/user_info_view.test.tsx | 253 ++++++++++++- .../_components/view_users/user_info_view.tsx | 101 +++++- .../src/components/networking.tsx | 1 + .../permissions/MCPServerPermissions.tsx | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 14 files changed, 1451 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index a27d6b92843..423cda5eea2 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 @@ -1,6 +1,6 @@ import re from datetime import datetime, timezone -from typing import Dict, List, Optional, Set, Tuple, cast +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Set, Tuple, cast from fastapi import HTTPException from starlette.datastructures import Headers @@ -30,6 +30,7 @@ ) from litellm.proxy._types import ( UI_TEAM_ID, + LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, ProxyException, SpecialHeaders, @@ -43,13 +44,27 @@ user_api_key_auth, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl +from litellm.proxy.common_utils.user_api_key_cache import ( + USER_NO_MCP_PERMISSION_SENTINEL, + get_management_object_ttl, + user_object_permission_id_cache_key, +) from litellm.repositories.table_repositories import ( AgentsRepository, MCPServerRepository, ) +from litellm.repositories.user_repository import UserRepository from litellm.types.mcp_server.mcp_server_manager import MCPServer +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +def _as_list(values: Sequence[str] | None) -> list[str] | None: # mutable-ok: resolver returns a list + """Widen a read-only allowlist back to the mutable list the resolver's own contract returns, + preserving the ``None`` that means "no restriction".""" + return None if values is None else list(values) + def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: Optional[List[str]] = None) -> Optional[List[str]]: """Resolve the single MCP server name a cold-start passthrough bypass may @@ -1408,6 +1423,15 @@ async def get_allowed_mcp_servers( f"Applied agent intersection filter. Final allowed servers: {allowed_mcp_servers}" ) + ######################################################### + # Apply the internal user's own ceiling (the entitlement attached to the human) + ######################################################### + capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling( + allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source + ) + allowed_mcp_servers = list(capped) + has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts + ######################################################### # Apply org-level ceiling if org_id is set ######################################################### @@ -1831,6 +1855,12 @@ async def get_allowed_tools_for_server( # No team restrictions → use key restrictions allowed_tools = cast(List[str], key_tools) + allowed_tools = _as_list( + await MCPRequestHandler._apply_user_tool_ceiling( + allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source + ) + ) + return await MCPRequestHandler._apply_agent_and_org_tool_ceilings( allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source ) @@ -2376,6 +2406,203 @@ async def _get_allowed_mcp_servers_for_end_user( verbose_logger.warning(f"Failed to get allowed MCP servers for end_user: {str(e)}") return [] + @staticmethod + async def _get_user_object_permission( + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> LiteLLM_ObjectPermissionTable | None: + """The internal user's OWN object_permission: the entitlement attached to the HUMAN rather + than to the credential they authenticated with. + + A key's object_permission is the credential's scope and a team's is the group's; this one + answers "which MCP servers and tools is this person entitled to", independent of how many keys + they hold. Caches the ``user_id -> object_permission_id`` mapping (with a sentinel for "no + entitlement") exactly as the agent path does, then reuses the shared ``object_permission_id`` + cache, so a warm request reads no rows. + + ``None`` means the human places NO ceiling: no user row, or a row naming no permission. The + two fault classes are deliberately NOT collapsed into that: a user row we cannot read leaves + us unable to say whether they are entitled at all, which is exactly the state before this + level existed, so it places no ceiling; a row that NAMES a permission we cannot read is a + KNOWN entitlement with unknown contents, so it raises and the caller denies. + """ + from litellm.proxy.auth.auth_checks import get_object_permission + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if not user_api_key_auth or not user_api_key_auth.user_id: + return None + + if prisma_client is None: + verbose_logger.debug("prisma_client is None") + return None + + user_id = user_api_key_auth.user_id + object_permission_id = await MCPRequestHandler._user_object_permission_id(user_id, prisma_client) + if object_permission_id is None: + return None + + object_permission = await get_object_permission( + object_permission_id=object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if object_permission is None: + raise ValueError( + f"user {user_id!r} names object_permission_id {object_permission_id!r} which could not be loaded" + ) + return object_permission + + @staticmethod + async def _user_object_permission_id(user_id: str, prisma_client: "PrismaClient") -> str | None: + """The permission row this human's user row links to, or None when they link none. + + Caches the link (with a sentinel for "links none") so a human without an entitlement costs no + DB read per MCP request. Anything other than an id string is treated as a cache MISS rather + than carried into the permission lookup, and a read that fails answers None: not knowing + whether someone is entitled is the state that existed before this level, so it places no + ceiling. Only a link we DID resolve can make the caller deny. + """ + from litellm.proxy.proxy_server import user_api_key_cache + + cache_key = user_object_permission_id_cache_key(user_id) + try: + cached: object = await user_api_key_cache.async_get_cache(key=cache_key) + if cached == USER_NO_MCP_PERMISSION_SENTINEL: + return None + if isinstance(cached, str) and cached: + return cached + user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + linked: object = getattr(user_row, "object_permission_id", None) if user_row is not None else None + object_permission_id = linked if isinstance(linked, str) and linked else None + await user_api_key_cache.async_set_cache( + key=cache_key, + value=object_permission_id or USER_NO_MCP_PERMISSION_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) + return object_permission_id + except Exception as e: # noqa: BLE001 # unknown whether entitled at all: no ceiling, as before + verbose_logger.warning(f"MCP user entitlement: link for {user_id!r} unresolved, no ceiling: {str(e)}") + return None + + @staticmethod + async def _get_allowed_mcp_servers_for_user( + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> Sequence[str] | None: + """The MCP servers the internal user is entitled to, as server ids. + + ``[]`` means this human places no restriction (allow-all from this level); ``None`` means the + ceiling is UNRESOLVED, which the caller denies on. Servers named only under + ``mcp_tool_permissions`` count as entitled, exactly as they do for a key or a team, so + granting one tool never requires naming its server twice. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + try: + object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth) + if object_permissions is None: + return [] + + direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or []) + access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups( + object_permissions.mcp_access_groups or [] + ) + tool_perm_servers = list( + global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys() + ) + return list(set(direct_mcp_servers + access_group_servers + tool_perm_servers)) + except Exception as e: # noqa: BLE001 # any resolution fault is an unresolved ceiling, never "no ceiling" + verbose_logger.warning(f"Failed to get allowed MCP servers for user: {str(e)}") + return None + + @staticmethod + async def _apply_user_server_ceiling( + allowed_mcp_servers: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None = None, + *, + keyless_source: bool = False, + ) -> tuple[tuple[str, ...], bool]: + """Narrow a resolved server list by the internal user's own entitlement. + + Returns the capped list and whether this human restricted it at all; the caller needs the + second value because an org list may only CAP a lower-level restriction, never replace one, so + a user ceiling has to be visible to the org step. + + RAISES when the entitlement is known but unreadable, which the resolver's own handler turns + into deny-all. That is the point of the level: dropping a ceiling we know exists is exactly the + silent widening it is there to prevent. + """ + if keyless_source: + return tuple(allowed_mcp_servers), False + entitled = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth) + if entitled is None: + raise ValueError( + f"MCP user ceiling unresolvable for user_id=" + f"{user_api_key_auth.user_id if user_api_key_auth else None!r}" + ) + if not entitled: + return tuple(allowed_mcp_servers), False + capped = tuple(server for server in allowed_mcp_servers if server in set(entitled)) + verbose_logger.debug(f"Applied user ceiling filter. Final allowed servers: {capped}") + return capped, True + + @staticmethod + async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool: + """Whether this human's own entitlement bounds their MCP access at all. + + True when they are entitled to a specific set of servers, and also when that entitlement is + UNRESOLVED — a caller uses this to decide whether it may skip the resolver, and skipping it on + a transient fault would widen access. + """ + entitled_servers = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth) + return entitled_servers is None or len(entitled_servers) > 0 + + @staticmethod + async def _apply_user_tool_ceiling( + allowed_tools: Sequence[str] | None, + server_id: str, + user_api_key_auth: UserAPIKeyAuth | None = None, + *, + keyless_source: bool = False, + ) -> Sequence[str] | None: + """Narrow a key/team tool allowlist by the internal user's own tool entitlement. + + The human's entitlement can only ever narrow: a user naming tools on ``server_id`` intersects + (and becomes the allowlist when no lower level restricts), while a user naming none places no + restriction. Returns ``[]`` (deny every tool on this server) when the entitlement cannot be + resolved, because the caller's own except-handler treats a raise as allow-all for key auth. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + if keyless_source: + return allowed_tools + + try: + object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth) + except Exception as e: # noqa: BLE001 # an unresolved human entitlement must deny, not widen + verbose_logger.warning(f"MCP user tool ceiling unresolvable, denying tools on {server_id!r}: {str(e)}") + return [] + + if object_permissions is None or not object_permissions.mcp_tool_permissions: + return allowed_tools + + user_tools = global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).get( + server_id + ) + if user_tools is None: + return allowed_tools + if allowed_tools is None: + return list(user_tools) + return list(set(allowed_tools) & set(user_tools)) + # Sentinel stored in cache when an agent has no object_permission, so we # don't re-query the DB on every MCP request for that agent. _AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__" diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 82b820d8cd9..332fb5c322e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2418,6 +2418,11 @@ async def get_allowed_mcp_servers(self, user_api_key_auth: Optional[UserAPIKeyAu and not is_admitted_subject and _user_has_admin_view(user_api_key_auth) and not has_explicit_object_permission + # An entitlement attached to the HUMAN binds them whatever their role: it is the + # person's scope, not the credential's, so an admin role is not a waiver of it. An + # UNRESOLVED entitlement also skips the shortcut, so the resolver denies rather than + # handing over the whole registry on a transient fault. + and not await MCPRequestHandler._user_places_mcp_ceiling(user_api_key_auth) ): verbose_logger.debug("Admin user without explicit object_permission - returning all servers") return list(self.get_registry().keys()) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c6d4ee1120a..94b6e9f2e5f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2782,6 +2782,7 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): updated_at: Optional[datetime] = None sso_user_id: Optional[str] = None teams: List[str] = [] # Just team IDs, not full team objects + object_permission: LiteLLM_ObjectPermissionTable | None = None from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402 diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c46bc110ca8..1b5dc3d4cec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -74,6 +74,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, get_management_object_ttl, + object_permission_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -2609,7 +2610,7 @@ async def get_object_permission( raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys") # check if in cache - key = "object_permission_id:{}".format(object_permission_id) + key = object_permission_cache_key(object_permission_id) deserialized_perm = await user_api_key_cache.async_get_cache( key=key, model_type=LiteLLM_ObjectPermissionTable, diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 09921a3ac1d..dbb5b2c24d0 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -150,6 +150,28 @@ async def async_set_cache_pipeline( # type: ignore[override] return await super().async_set_cache_pipeline(cache_list=normalized, local_only=local_only, **kwargs) +#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row, +#: so a human without an entitlement costs no DB read per request. Lives beside the key builder +#: because it is part of the same cache protocol: a reader that knows the key must know this value. +USER_NO_MCP_PERMISSION_SENTINEL = "__user_no_mcp_permission__" + + +def user_object_permission_id_cache_key(user_id: str) -> str: + """Cache key for the ``user_id -> object_permission_id`` link. + + Lives here rather than next to either user because two modules own the two halves: the MCP auth + resolver writes it on read, and ``/user/update`` deletes it after changing the link. A key format + duplicated across those two drifts silently, and the failure is an entitlement change that never + takes effect. + """ + return f"user_object_permission_id:{user_id}" + + +def object_permission_cache_key(object_permission_id: str) -> str: + """Cache key ``get_object_permission`` stores a permission row under.""" + return f"object_permission_id:{object_permission_id}" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 1bd0a19bfb3..80d9ee21a44 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -45,6 +45,14 @@ generate_key_helper_fn, prepare_metadata_fields, ) +from litellm.proxy.common_utils.user_api_key_cache import ( + object_permission_cache_key, + user_object_permission_id_cache_key, +) +from litellm.proxy.management_helpers.object_permission_utils import ( + _set_object_permission, + handle_update_object_permission_common, +) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy, hash_password from litellm.repositories.organization_repository import OrganizationRepository @@ -401,7 +409,7 @@ async def new_user( - duration: Optional[str] - Duration for the key auto-created on `/user/new`. Default is None. - key_alias: Optional[str] - Alias for the key auto-created on `/user/new`. Default is None. - sso_user_id: Optional[str] - The id of the user in the SSO provider. - - object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1"], "mcp_servers": ["github"], "mcp_tool_permissions": {"github": ["list_issues"]}}. The MCP grants act as a ceiling on every key this user holds. IF null or {} then no object permission. - prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts. - organizations: List[str] - List of organization id's the user is a member of - budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -466,6 +474,10 @@ async def new_user( data_json = data.json() # type: ignore data_json = _update_internal_new_user_params(data_json, data) + # Persist the requested grants as their own row and link it, mirroring key/team creation. + # generate_key_helper_fn only forwards object_permission_id, so without this the entitlement + # the caller sent would be dropped on the floor. + data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client) _hash_password_in_dict(data_json) teams = data.teams if teams is None: @@ -852,9 +864,12 @@ async def _check_user_info_v2_access( if prisma_client is None: return None - # Helper: fetch the target user row (reused across branches) + # Helper: fetch the target user row (reused across branches). object_permission is included so + # callers can read the user's MCP/vector-store entitlements without a second round trip. async def _fetch_target_user(): - return await UserRepository(prisma_client).table.find_unique(where={"user_id": target_user_id}) + return await UserRepository(prisma_client).table.find_unique( + where={"user_id": target_user_id}, include={"object_permission": True} + ) # Rule 1: Proxy admins — fetch and return the target row directly if _user_has_admin_view(user_api_key_dict): @@ -972,6 +987,7 @@ async def user_info_v2( updated_at=user_data.get("updated_at"), sso_user_id=user_data.get("sso_user_id"), teams=user_data.get("teams") or [], + object_permission=user_data.get("object_permission"), ) except Exception as e: verbose_proxy_logger.exception( @@ -1207,6 +1223,48 @@ async def _invalidate_user_spend_counter_if_changed( await _invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}") +def _clears_object_permission(user_request: UpdateUserRequest) -> bool: + """Whether the caller explicitly asked to remove this user's object_permission. + + Distinguishes "sent nothing" from "sent an empty grant set". Only the latter clears; an omitted + field must leave an existing entitlement alone. + """ + if "object_permission" not in (user_request.fields_set() if hasattr(user_request, "fields_set") else set()): + return False + sent = user_request.object_permission + return sent is None or not sent.model_dump(exclude_unset=True, exclude_none=True) + + +async def _invalidate_cached_user_entitlement(user_id: str | None, object_permission_ids: tuple[str, ...]) -> None: + """Drop the cache entries an entitlement change makes stale. + + All three kinds are needed: a permission row is cached under its own id (so re-reading the same + link still yields the OLD grants), the ``user_id -> object_permission_id`` link is cached + separately (so a user who previously had NO entitlement keeps its "none" sentinel), and the user + row itself is cached whole. Leaving any behind means an admin revoking a tool keeps serving it + until the management-object TTL expires. + + Both the outgoing and incoming permission ids are passed, because a clear leaves no incoming id + at all and an upsert may mint a new row; invalidating only one of the two leaves the other's + grants live. + + Each deletion is isolated: one that fails must not skip the others, or a single unreachable key + would silently leave the rest of a revocation in place. Best-effort overall, exactly as the caches + are everywhere else, since one we cannot clear still expires on its own. + """ + from litellm.proxy.proxy_server import user_api_key_cache + + keys = ( + *(object_permission_cache_key(permission_id) for permission_id in dict.fromkeys(object_permission_ids)), + *((user_object_permission_id_cache_key(user_id), user_id) if user_id is not None else ()), + ) + for key in keys: + try: + await user_api_key_cache.async_delete_cache(key=key) + except Exception as e: # noqa: BLE001 # a cache we cannot clear still expires; never fail the write + verbose_proxy_logger.warning(f"Failed to invalidate cached entitlement key {key!r}: {str(e)}") + + async def _update_single_user_helper( user_request: UpdateUserRequest, user_api_key_dict: UserAPIKeyAuth, @@ -1259,9 +1317,15 @@ async def _update_single_user_helper( ) _is_self_update = _target_user_id is not None and user_api_key_dict.user_id == _target_user_id if _is_self_update and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: - _protected_fields = ("max_budget", "soft_budget", "spend") + # object_permission is a CEILING on what this human may reach, so a self-write is an + # escalation path: sending an empty grant list means "no restriction" and would lift a + # restriction an admin placed on them. Checked against the fields the caller actually SENT, + # because `_update_internal_user_params` drops empty values, and `object_permission: {}` is + # precisely the clear-my-own-ceiling case this must refuse. + _sent_fields = user_request.fields_set() if hasattr(user_request, "fields_set") else set() + _protected_fields = ("max_budget", "soft_budget", "spend", "object_permission") for _field in _protected_fields: - if _field in non_default_values: + if _field in non_default_values or _field in _sent_fields: raise HTTPException( status_code=403, detail={ @@ -1282,6 +1346,22 @@ async def _update_single_user_helper( # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(non_default_values.get("spend")) + # Upsert the grants into their own row and link it, mirroring /key/update and /team/update. + # This also removes object_permission from the payload, which is not a column on the user table. + if "object_permission" in non_default_values: + object_permission_id = await handle_update_object_permission_common( + data_json=non_default_values, + existing_object_permission_id=getattr(existing_user_row, "object_permission_id", None), + prisma_client=prisma_client, + ) + if object_permission_id is not None: + non_default_values["object_permission_id"] = object_permission_id + elif _clears_object_permission(user_request): + # An explicit `{}` or null means "no object permission", which the merge-based upsert cannot + # express: merging an empty grant set over the existing row leaves every grant in place. So + # the link is dropped instead, which is what makes the documented clear actually clear. + non_default_values["object_permission_id"] = None + # Perform the update response: dict[str, Any] | None = None @@ -1326,6 +1406,19 @@ async def _update_single_user_helper( await _invalidate_user_spend_counter_if_changed(non_default_values) + if "object_permission_id" in non_default_values: + await _invalidate_cached_user_entitlement( + user_id=non_default_values.get("user_id"), + object_permission_ids=tuple( + permission_id + for permission_id in ( + getattr(existing_user_row, "object_permission_id", None), + non_default_values.get("object_permission_id"), + ) + if isinstance(permission_id, str) + ), + ) + if response is None: raise HTTPException( status_code=400, @@ -1407,7 +1500,7 @@ async def user_update( - team_id: Optional[str] - [DEPRECATED PARAM] The team id of the user. Default is None. - duration: Optional[str] - [NOT IMPLEMENTED]. - key_alias: Optional[str] - [NOT IMPLEMENTED]. - - object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1"], "mcp_servers": ["github"], "mcp_tool_permissions": {"github": ["list_issues"]}}. The MCP grants act as a ceiling on every key this user holds. IF null or {} then no object permission. - prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts. - budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index b3c0dcd1681..4be2bb053ef 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7645,3 +7645,346 @@ async def test_scrub_admitted_drops_authorization_but_keeps_injected_upstream_to assert oauth2 is None assert "authorization" not in {k.lower() for k in raw} assert per_server == {"github": {"Authorization": "Bearer gh_injected_upstream"}} + + +# --------------------------------------------------------------------------- +# Internal-user (human) MCP entitlement tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestUserMCPEntitlement: + """The entitlement attached to the HUMAN, read at both list time and tool-call time. + + A key's object_permission scopes the credential and a team's scopes the group; the user's own + scopes the person, so it must cap every key they hold and every tool those keys may invoke. + """ + + def _auth(self, user_id: str = "human-1", **kwargs) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_id=user_id, **kwargs) + + def _perm(self, *, servers=None, access_groups=None, tool_permissions=None) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable( + object_permission_id="perm-human-1", + mcp_servers=servers if servers is not None else [], + mcp_access_groups=access_groups if access_groups is not None else [], + mcp_tool_permissions=tool_permissions, + ) + + @contextlib.contextmanager + def _entitled(self, perm): + """Patch the human's entitlement lookup. ``perm`` may be a permission row, None, or an + exception instance to raise (an entitlement that cannot be resolved).""" + side_effect = perm if isinstance(perm, Exception) else None + with patch.object( + MCPRequestHandler, + "_get_user_object_permission", + new_callable=AsyncMock, + return_value=None if side_effect else perm, + side_effect=side_effect, + ) as patched: + yield patched + + @contextlib.contextmanager + def _key_and_team_servers(self, key_servers, team_servers): + with ( + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_key", + new_callable=AsyncMock, + return_value=key_servers, + ), + patch.object( + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + new_callable=AsyncMock, + return_value=team_servers, + ), + patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ), + patch.object( + MCPRequestHandler, + "_get_key_access_group_mcp_server_extras", + new_callable=AsyncMock, + return_value=[], + ), + ): + yield + + async def test_entitlement_caps_the_servers_the_key_reaches(self): + """The key grants two servers; the human is entitled to one, so only that one resolves.""" + with self._key_and_team_servers(["srv-a", "srv-b"], []): + with self._entitled(self._perm(servers=["srv-a"])): + result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth()) + assert result == ["srv-a"] + + async def test_entitlement_never_widens_the_key(self): + """A human entitled to a server their key does not grant still cannot reach it: the level is a + ceiling, so it intersects rather than unions.""" + with self._key_and_team_servers(["srv-a"], []): + with self._entitled(self._perm(servers=["srv-a", "srv-elsewhere"])): + result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth()) + assert result == ["srv-a"] + + async def test_no_entitlement_places_no_ceiling(self): + """A human with no entitlement row leaves the key/team result untouched.""" + with self._key_and_team_servers(["srv-a", "srv-b"], []): + with self._entitled(None): + result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth()) + assert sorted(result) == ["srv-a", "srv-b"] + + async def test_unresolvable_entitlement_denies_every_server(self): + """A KNOWN entitlement whose contents cannot be read must deny, not fall back to the key's + wider scope.""" + with self._key_and_team_servers(["srv-a", "srv-b"], []): + with self._entitled(ValueError("permission row unreadable")): + result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth()) + assert result == [] + + async def test_entitlement_caps_the_tools_the_key_reaches(self): + """Tool-level: the key allows three tools on the server, the human is entitled to one.""" + key_perm = self._perm(tool_permissions={"srv-a": ["read", "write", "delete"]}) + with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_perm): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None + ): + with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})): + result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth()) + assert result == ["read"] + + async def test_entitlement_alone_restricts_tools_on_an_otherwise_unrestricted_key(self): + """An unrestricted key (no tool permissions of its own) is still bound by the human's tools.""" + with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None + ): + with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})): + result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth()) + assert result == ["read"] + + async def test_entitlement_on_another_server_does_not_restrict_this_one(self): + """Tool grants are per server: naming tools on srv-b places no bound on srv-a.""" + with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None + ): + with self._entitled(self._perm(tool_permissions={"srv-b": ["read"]})): + result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth()) + assert result is None + + async def test_unresolvable_entitlement_denies_every_tool(self): + """Fail closed on the tool axis too. The caller's own except-handler treats a raise as + allow-all for key auth, so the ceiling must return the empty allowlist itself.""" + with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None + ): + with self._entitled(ValueError("permission row unreadable")): + result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth()) + assert result == [] + + async def test_tool_call_is_rejected_at_call_time(self): + """The end-to-end contract: a tool the human is not entitled to is refused when INVOKED, not + merely hidden from the advertised list.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv-a", + name="srv-a", + server_name="srv-a", + url="https://srv-a.example.com", + transport=MCPTransport.http, + ) + with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None + ): + with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})): + await global_mcp_server_manager.check_tool_permission_for_key_team( + tool_name="read", server=server, user_api_key_auth=self._auth() + ) + with pytest.raises(HTTPException) as exc: + await global_mcp_server_manager.check_tool_permission_for_key_team( + tool_name="delete", server=server, user_api_key_auth=self._auth() + ) + assert exc.value.status_code == 403 + + async def test_keyless_admitted_source_is_not_capped_by_the_user_level(self): + """A gateway-admitted human resolves as a UNION over their own grants plus their teams', and + their own grants ARE the user source there. Re-applying them as a ceiling per source would + make one team's narrower scope silently bound another's, so the level is skipped.""" + with self._key_and_team_servers(["srv-a", "srv-b"], []): + with self._entitled(self._perm(servers=["srv-a"])) as lookup: + result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth(), keyless_source=True) + assert sorted(result) == ["srv-a", "srv-b"] + lookup.assert_not_awaited() + + async def test_keyless_admitted_source_tools_are_not_capped_by_the_user_level(self): + with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None): + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None + ): + with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})) as lookup: + result = await MCPRequestHandler.get_allowed_tools_for_server( + "srv-a", self._auth(), keyless_source=True + ) + assert result is None + lookup.assert_not_awaited() + + async def test_servers_named_only_under_tool_permissions_are_entitled(self): + """Granting one tool on a server entitles the human to that server, so an admin never has to + name it twice.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + global_mcp_server_manager.registry["srv-a"] = MCPServer( + server_id="srv-a", + name="srv-a", + server_name="srv-a", + url="https://srv-a.example.com", + transport=MCPTransport.http, + ) + try: + with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})): + with patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_user(self._auth()) + finally: + global_mcp_server_manager.registry.pop("srv-a", None) + assert result == ["srv-a"] + + async def test_places_ceiling_is_true_when_unresolvable(self): + """``_user_places_mcp_ceiling`` gates the admin shortcut that hands over the whole registry, so + an entitlement it cannot resolve must still count as a ceiling.""" + with self._entitled(ValueError("boom")): + assert await MCPRequestHandler._user_places_mcp_ceiling(self._auth()) is True + with self._entitled(None): + assert await MCPRequestHandler._user_places_mcp_ceiling(self._auth()) is False + + +@pytest.mark.asyncio +class TestGetUserObjectPermission: + """Resolution of the ``user_id -> object_permission_id -> grants`` chain.""" + + def _prisma_with_user(self, user_row): + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + return prisma_client + + async def test_resolves_through_the_shared_permission_cache(self): + from litellm.caching.dual_cache import DualCache + + user_row = MagicMock() + user_row.object_permission_id = "perm-1" + prisma_client = self._prisma_with_user(user_row) + auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-shared") + expected = MagicMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_object_permission", + new_callable=AsyncMock, + return_value=expected, + ) as mock_get_perm, + ): + assert await MCPRequestHandler._get_user_object_permission(auth) is expected + assert mock_get_perm.await_args.kwargs["object_permission_id"] == "perm-1" + + # The user_id -> object_permission_id link is cached, so the user row is read once. + prisma_client.db.litellm_usertable.find_unique.reset_mock() + await MCPRequestHandler._get_user_object_permission(auth) + prisma_client.db.litellm_usertable.find_unique.assert_not_called() + + async def test_caches_a_sentinel_for_a_human_with_no_entitlement(self): + """A human without an entitlement is the common case and must cost no DB read per request.""" + from litellm.caching.dual_cache import DualCache + + user_row = MagicMock() + user_row.object_permission_id = None + prisma_client = self._prisma_with_user(user_row) + auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-no-perm") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch("litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock) as mock_get_perm, + ): + assert await MCPRequestHandler._get_user_object_permission(auth) is None + assert await MCPRequestHandler._get_user_object_permission(auth) is None + mock_get_perm.assert_not_awaited() + prisma_client.db.litellm_usertable.find_unique.assert_awaited_once() + + async def test_missing_user_row_places_no_ceiling(self): + """Whether this human is entitled at all is unknown when their row is absent, which is the + state before the level existed, so it must not deny.""" + from litellm.caching.dual_cache import DualCache + + prisma_client = self._prisma_with_user(None) + auth = UserAPIKeyAuth(api_key="sk-test", user_id="ghost") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ): + assert await MCPRequestHandler._get_user_object_permission(auth) is None + + async def test_unreadable_user_row_places_no_ceiling(self): + from litellm.caching.dual_cache import DualCache + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=Exception("db down")) + auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-db-down") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + ): + assert await MCPRequestHandler._get_user_object_permission(auth) is None + + async def test_named_but_unreadable_permission_raises(self): + """A KNOWN entitlement with unknown contents is indeterminate: it must surface so the callers + can deny rather than serve the wider key scope.""" + from litellm.caching.dual_cache import DualCache + + user_row = MagicMock() + user_row.object_permission_id = "perm-gone" + prisma_client = self._prisma_with_user(user_row) + auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-dangling") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_object_permission", + new_callable=AsyncMock, + return_value=None, + ), + ): + with pytest.raises(ValueError): + await MCPRequestHandler._get_user_object_permission(auth) + + async def test_no_user_id_places_no_ceiling(self): + assert await MCPRequestHandler._get_user_object_permission(UserAPIKeyAuth(api_key="sk-test")) is None + assert await MCPRequestHandler._get_user_object_permission(None) is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 8de42ca89da..a37f7ca764d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2893,6 +2893,7 @@ async def mock_find_unique(*args, **kwargs): "updated_at", "sso_user_id", "teams", + "object_permission", } assert set(response_dict.keys()) == expected_fields @@ -3702,3 +3703,332 @@ async def test_get_user_info_for_proxy_admin_validates_keys_and_teams(): returned_key = result.keys[0] assert returned_key["team_id"] == "team-a" assert returned_key["models"] == [] + + +def _object_permission_mocks(mocker, existing_object_permission_id=None): + """Prisma double whose user row optionally already links a permission row.""" + mock_prisma_client = mocker.MagicMock() + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = { + "user_id": "target-user", + "object_permission_id": existing_object_permission_id, + } + existing_user.user_id = "target-user" + existing_user.object_permission_id = existing_object_permission_id + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=existing_user + ) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = mocker.AsyncMock( + return_value=SimpleNamespace(object_permission_id="perm-new") + ) + mock_prisma_client.update_data = mocker.AsyncMock( + return_value={"user_id": "target-user"} + ) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + mocker.patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + new=mocker.AsyncMock(), + ) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_user_update_persists_mcp_entitlement_and_links_it(mocker): + """/user/update documents an object_permission param; it must actually be stored. + + The grants live in their own table, so the endpoint has to upsert them and hand the user row + only the resulting object_permission_id. Passing object_permission through to the user update + would not even be a column. + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = _object_permission_mocks(mocker) + cache = mocker.MagicMock() + cache.async_delete_cache = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id="target-user", + object_permission={ + "mcp_servers": ["github"], + "mcp_tool_permissions": {"github": ["list_issues"]}, + }, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + upsert_kwargs = mock_prisma_client.db.litellm_objectpermissiontable.upsert.call_args.kwargs + created = upsert_kwargs["data"]["create"] + assert created["mcp_servers"] == ["github"] + assert json.loads(created["mcp_tool_permissions"]) == {"github": ["list_issues"]} + + written = mock_prisma_client.update_data.call_args.kwargs["data"] + assert written["object_permission_id"] == "perm-new" + assert "object_permission" not in written + + +@pytest.mark.asyncio +async def test_user_update_invalidates_the_cached_entitlement(mocker): + """An admin revoking a tool must take effect now, not at the end of the cache TTL. + + Three entries go stale: the permission row (keyed by its own id), the user -> permission link + (which carries a "no entitlement" sentinel), and the cached user row. + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + _object_permission_mocks(mocker) + cache = mocker.MagicMock() + cache.async_delete_cache = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id="target-user", + object_permission={"mcp_tool_permissions": {"github": []}}, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} + assert deleted == { + "object_permission_id:perm-new", + "user_object_permission_id:target-user", + "target-user", + } + + +@pytest.mark.asyncio +async def test_admin_can_clear_a_users_mcp_entitlement(mocker): + """An explicit empty object_permission means "no object permission", so it must unlink. + + The merge-based upsert cannot express this: merging an empty grant set over the existing row + leaves every grant in place, and the empty-value filter drops the field before the upsert runs, + so without the explicit clear path the documented operation silently returns success unchanged. + + A clear also leaves no incoming permission id, so invalidation keyed off one would skip it and + the gateway would keep enforcing the cleared grants until the cache expired. + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = _object_permission_mocks(mocker, "perm-existing") + cache = mocker.MagicMock() + cache.async_delete_cache = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + + await _update_single_user_helper( + user_request=UpdateUserRequest(user_id="target-user", object_permission={}), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + written = mock_prisma_client.update_data.call_args.kwargs["data"] + assert written["object_permission_id"] is None + assert "object_permission" not in written + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called() + + deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} + assert deleted == { + "object_permission_id:perm-existing", + "user_object_permission_id:target-user", + "target-user", + } + + +@pytest.mark.asyncio +async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mocker): + """An upsert can mint a new permission row, which leaves the outgoing one cached under its id. + + Only the link cache knows the user moved; the old row's own entry still holds the pre-update + grants, so anything still resolving that id keeps reading them. + """ + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + _object_permission_mocks(mocker, "perm-existing") + cache = mocker.MagicMock() + cache.async_delete_cache = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id="target-user", + object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} + assert deleted == { + "object_permission_id:perm-existing", + "object_permission_id:perm-new", + "user_object_permission_id:target-user", + "target-user", + } + + +@pytest.mark.asyncio +async def test_non_admin_cannot_clear_their_own_mcp_entitlement(mocker): + """The empty-value filter drops `object_permission: {}` before the guard saw it, so a non-admin + could clear the very ceiling an admin placed on them. The guard reads the fields the caller SENT. + """ + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = _object_permission_mocks(mocker, "perm-existing") + cache = mocker.MagicMock() + cache.async_delete_cache = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + + with pytest.raises(HTTPException) as exc: + await _update_single_user_helper( + user_request=UpdateUserRequest(user_id="target-user", object_permission={}), + user_api_key_dict=UserAPIKeyAuth( + user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER + ), + ) + + assert exc.value.status_code == 403 + mock_prisma_client.update_data.assert_not_called() + + +@pytest.mark.asyncio +async def test_non_admin_cannot_rewrite_their_own_mcp_entitlement(mocker): + """The entitlement bounds the human, so a self-write is an escalation path: an empty grant list + means "no restriction" and would lift a ceiling the admin placed on them.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + mock_prisma_client = _object_permission_mocks(mocker, "perm-existing") + cache = mocker.MagicMock() + cache.async_delete_cache = mocker.AsyncMock() + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) + + with pytest.raises(HTTPException) as exc: + await _update_single_user_helper( + user_request=UpdateUserRequest( + user_id="target-user", + object_permission={"mcp_servers": [], "mcp_tool_permissions": {}}, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER + ), + ) + + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called() + mock_prisma_client.update_data.assert_not_called() + + +@pytest.mark.asyncio +async def test_new_user_persists_the_requested_mcp_entitlement(mocker): + """generate_key_helper_fn only forwards object_permission_id, so /user/new has to create the + grants row itself; otherwise the entitlement the admin sent is silently dropped.""" + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.create = mocker.AsyncMock( + return_value=SimpleNamespace(object_permission_id="perm-created") + ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_usertable.count = mocker.AsyncMock(return_value=0) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.check_if_default_team_set", + return_value=None, + ) + mock_generate = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", + new=mocker.AsyncMock( + return_value={"user_id": "new-human", "token": "sk-x", "expires": None} + ), + ) + mocker.patch( + "litellm.proxy.hooks.user_management_event_hooks.UserManagementEventHooks.async_user_created_hook", + new=mocker.AsyncMock(), + ) + + await new_user( + data=NewUserRequest( + user_id="new-human", + object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + created = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert json.loads(created["mcp_tool_permissions"]) == {"github": ["list_issues"]} + forwarded = mock_generate.call_args.kwargs + assert forwarded["object_permission_id"] == "perm-created" + assert "object_permission" not in forwarded + + +@pytest.mark.asyncio +async def test_user_info_v2_returns_the_mcp_entitlement(mocker): + """The admin UI reads the current entitlement off this endpoint, so the grants have to come back + with the user row rather than only their id.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 + + user_row = SimpleNamespace( + object_permission=SimpleNamespace( + object_permission_id="perm-1", + mcp_servers=["github"], + mcp_access_groups=[], + mcp_tool_permissions={"github": ["list_issues"]}, + ), + ) + user_row.model_dump = lambda: { + "user_id": "human-1", + "object_permission": { + "object_permission_id": "perm-1", + "mcp_servers": ["github"], + "mcp_access_groups": [], + "mcp_tool_permissions": {"github": ["list_issues"]}, + }, + } + mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.MagicMock()) + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_user_info_v2_access", + new=mocker.AsyncMock(return_value=user_row), + ) + + response = await user_info_v2( + request=SimpleNamespace(query_params={}), + user_id="human-1", + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + assert response.object_permission is not None + assert response.object_permission.mcp_servers == ["github"] + assert response.object_permission.mcp_tool_permissions == { + "github": ["list_issues"] + } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx index 54392adf885..c83a8e48e43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/user_edit_view.tsx @@ -1,11 +1,14 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, SelectItem, TextInput, Textarea } from "@tremor/react"; -import { Checkbox, Form, Select, Tooltip } from "antd"; +import { Checkbox, Form, Input, Select, Tooltip } from "antd"; import React, { useState } from "react"; import { all_admin_roles } from "@/utils/roles"; import BudgetDurationDropdown from "@/components/common_components/budget_duration_dropdown"; import { getModelDisplayName } from "@/components/key_team_helpers/fetch_available_models_team_key"; import NumericalInput from "@/components/shared/numerical_input"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; +import type { ObjectPermission } from "@/components/object_permission_types"; interface UserEditViewProps { userData: any; @@ -18,8 +21,18 @@ interface UserEditViewProps { userModels: string[]; possibleUIRoles: Record> | null; isBulkEdit?: boolean; + objectPermission?: ObjectPermission | null; } +const buildMcpFieldValues = (objectPermission: ObjectPermission | null | undefined) => ({ + mcp_servers_and_groups: { + servers: objectPermission?.mcp_servers ?? [], + accessGroups: objectPermission?.mcp_access_groups ?? [], + toolsets: objectPermission?.mcp_toolsets ?? [], + }, + mcp_tool_permissions: objectPermission?.mcp_tool_permissions ?? {}, +}); + export function UserEditView({ userData, onCancel, @@ -31,9 +44,11 @@ export function UserEditView({ userModels, possibleUIRoles, isBulkEdit = false, + objectPermission, }: UserEditViewProps) { const [form] = Form.useForm(); const [unlimitedBudget, setUnlimitedBudget] = useState(false); + const canEditMcpPermissions = !isBulkEdit && all_admin_roles.includes(userRole || ""); // Set initial form values React.useEffect(() => { @@ -50,8 +65,9 @@ export function UserEditView({ max_budget: isUnlimited ? "" : maxBudget, budget_duration: userData.user_info?.budget_duration, metadata: userData.user_info?.metadata ? JSON.stringify(userData.user_info.metadata, null, 2) : undefined, + ...(canEditMcpPermissions ? buildMcpFieldValues(objectPermission) : {}), }); - }, [userData, form]); + }, [userData, objectPermission, canEditMcpPermissions, form]); const handleUnlimitedBudgetChange = (e: any) => { const checked = e.target.checked; @@ -186,6 +202,52 @@ export function UserEditView({