diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 5b56133f1ce..672c38d16a7 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -63,6 +63,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, attach_object_permission_to_dict, handle_update_object_permission_common, + validate_mcp_object_permission_for_key, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, @@ -620,6 +621,12 @@ async def _common_key_generation_helper( # noqa: PLR0915 data_json.pop("tags") + await validate_mcp_object_permission_for_key( + user_api_key_dict=user_api_key_dict, + object_permission=data_json.get("object_permission"), + prisma_client=prisma_client, + ) + data_json = await _set_object_permission( data_json=data_json, prisma_client=prisma_client, @@ -1927,6 +1934,13 @@ async def update_key_fn( # Set Management Endpoint Metadata Fields + if "object_permission" in data_json: + await validate_mcp_object_permission_for_key( + user_api_key_dict=user_api_key_dict, + object_permission=data_json.get("object_permission"), + prisma_client=prisma_client, + ) + non_default_values = await prepare_key_update_data( data=data, existing_key_row=existing_key_row ) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 8c4d4e7937e..f23dd75a753 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -455,11 +455,15 @@ if MCP_AVAILABLE: user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Get all available MCP access groups from the database AND config + Get MCP access groups available to the user. Non-admins only see groups + they have access to via their teams; admins see all groups. """ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.management_helpers.object_permission_utils import ( + get_allowed_mcp_access_groups_for_user, + ) from litellm.proxy.proxy_server import prisma_client access_groups = set() @@ -482,6 +486,14 @@ if MCP_AVAILABLE: except Exception as e: verbose_proxy_logger.debug(f"Error getting MCP access groups: {e}") + # Filter for non-admins: only return groups the user has access to + allowed_groups = await get_allowed_mcp_access_groups_for_user( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + if allowed_groups is not None: + access_groups = access_groups & allowed_groups + # Convert to sorted list access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 9670cdf330a..768c27b7483 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -5,12 +5,14 @@ organizations, teams, and keys. import json from litellm._uuid import uuid -from typing import Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Union + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - async def attach_object_permission_to_dict( @@ -177,4 +179,155 @@ async def _set_object_permission( data_json["object_permission_id"] = created_permission.object_permission_id data_json.pop("object_permission") - return data_json \ No newline at end of file + return data_json + + +async def get_allowed_mcp_access_groups_for_user( + user_api_key_dict: "UserAPIKeyAuth", + prisma_client: Optional[PrismaClient], +) -> Optional[set]: + """ + Return the set of MCP access group IDs the user can assign (via teams or key). + Returns None if user is admin (can assign any) or if MCP modules are unavailable. + """ + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + + if _user_has_admin_view(user_api_key_dict): + return None # Admin can assign any + + try: + from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + build_effective_auth_contexts, + ) + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import user_api_key_cache, proxy_logging_obj + except ImportError: + return set() # No MCP - user has no access + + if prisma_client is None or user_api_key_cache is None: + return set() + + allowed_access_group_ids: set = set() + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + + for auth_context in auth_contexts: + if auth_context.team_id: + team_obj = await get_team_object( + team_id=auth_context.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=getattr( + user_api_key_dict, "parent_otel_span", None + ), + proxy_logging_obj=proxy_logging_obj, + ) + if team_obj and team_obj.object_permission: + groups = team_obj.object_permission.mcp_access_groups or [] + allowed_access_group_ids.update(groups) + + if user_api_key_dict.object_permission: + key_groups = user_api_key_dict.object_permission.mcp_access_groups or [] + allowed_access_group_ids.update(key_groups) + + return allowed_access_group_ids + + +async def validate_mcp_object_permission_for_key( + user_api_key_dict: "UserAPIKeyAuth", + object_permission: Optional[Union[Dict[str, Any], Any]], + prisma_client: Optional[PrismaClient], +) -> None: + """ + Validate that a non-admin user can only assign MCP servers and access groups + they have access to (via their teams or key). With view_all mode, users see + all servers but must not be able to assign servers/groups they lack access to. + + Raises: + HTTPException: 403 if user tries to assign MCP servers or access groups + they do not have access to. + """ + from fastapi import HTTPException, status + + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + + if object_permission is None: + return + + # Admins can assign any MCP servers/groups + if _user_has_admin_view(user_api_key_dict): + return + + # Extract mcp_servers and mcp_access_groups from object_permission + mcp_servers: list = [] + mcp_access_groups: list = [] + if isinstance(object_permission, dict): + mcp_servers = object_permission.get("mcp_servers") or [] + mcp_access_groups = object_permission.get("mcp_access_groups") or [] + else: + mcp_servers = getattr(object_permission, "mcp_servers", None) or [] + mcp_access_groups = getattr(object_permission, "mcp_access_groups", None) or [] + + if not mcp_servers and not mcp_access_groups: + return + + try: + from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + build_effective_auth_contexts, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import user_api_key_cache + except ImportError: + verbose_proxy_logger.debug( + "MCP modules not available, skipping MCP object permission validation" + ) + return + + if prisma_client is None or user_api_key_cache is None: + return + + allowed_server_ids: set = set() + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + for auth_context in auth_contexts: + server_ids = await global_mcp_server_manager.get_allowed_mcp_servers( + auth_context + ) + allowed_server_ids.update(server_ids) + + allowed_access_group_ids = await get_allowed_mcp_access_groups_for_user( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + if allowed_access_group_ids is None: + allowed_access_group_ids = set() + + # Validate requested mcp_servers + disallowed_servers = [ + s for s in mcp_servers if s not in allowed_server_ids + ] + if disallowed_servers: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"You do not have access to assign the following MCP servers to this key: {disallowed_servers}. " + "You can only assign MCP servers that your teams have access to." + ) + }, + ) + + # Validate requested mcp_access_groups + disallowed_groups = [ + g for g in mcp_access_groups if g not in allowed_access_group_ids + ] + if disallowed_groups: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"You do not have access to assign the following MCP access groups to this key: {disallowed_groups}. " + "You can only assign MCP access groups that your teams have access to." + ) + }, + ) \ No newline at end of file diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 7565e901ecd..cb517dd45b3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -642,6 +642,65 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): assert created_permission_data["mcp_servers"] == ["server_1"] +@pytest.mark.asyncio +async def test_validate_mcp_object_permission_rejects_unauthorized_servers(): + """ + With view_all mode, non-admin users see all MCP servers but must not be able to + assign servers they lack access to. validate_mcp_object_permission_for_key + should raise 403 when a user with no MCP access tries to assign mcp_servers. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_helpers.object_permission_utils import ( + validate_mcp_object_permission_for_key, + ) + + # Non-admin user with no team - has no MCP access + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-user-key", + user_id="test_user", + team_id=None, + ) + + mock_prisma = MagicMock() + mock_user_cache = MagicMock() + + # build_effective_auth_contexts returns single context (no teams) + # get_allowed_mcp_servers returns [] for user with no team/key MCP access + mock_manager = MagicMock() + mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.ui_session_utils.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user]), + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.proxy_server.user_api_key_cache", + mock_user_cache, + ), + ): + + with pytest.raises(HTTPException) as exc_info: + await validate_mcp_object_permission_for_key( + user_api_key_dict=mock_user, + object_permission={"mcp_servers": ["sensitive_server"]}, + prisma_client=mock_prisma, + ) + + assert exc_info.value.status_code == 403 + assert "sensitive_server" in str(exc_info.value.detail) + assert "do not have access" in str(exc_info.value.detail).lower() + + @pytest.mark.asyncio async def test_key_update_object_permissions_existing_permission(monkeypatch): """