mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
fixed mcp access group issue
This commit is contained in:
shivam
parent
88ccffccc8
commit
369062d6e0
4 changed files with 242 additions and 4 deletions
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
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."
|
||||
)
|
||||
},
|
||||
)
|
||||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue