diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index f3ec7da866c..c424f51774e 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -97,9 +97,15 @@ async def delete_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> Optional[MCPToolset]: - row = await prisma_client.db.litellm_mcptoolsettable.delete( - where={"toolset_id": toolset_id} - ) - if row is None: - return None + try: + row = await prisma_client.db.litellm_mcptoolsettable.delete( + where={"toolset_id": toolset_id} + ) + except Exception as e: + if ( + "RecordNotFoundError" in type(e).__name__ + or "record was not found" in str(e).lower() + ): + return None + raise return _toolset_from_row(row) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 44e567a90fd..e2a770f6b5d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -855,6 +855,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): mcp_servers: Optional[List[str]] = None mcp_access_groups: Optional[List[str]] = None mcp_tool_permissions: Optional[Dict[str, List[str]]] = None + mcp_toolsets: Optional[List[str]] = None vector_stores: Optional[List[str]] = None agents: Optional[List[str]] = None agent_access_groups: Optional[List[str]] = None diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 703d9360a3d..ed4d1bc36c0 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -380,30 +380,26 @@ async def validate_key_mcp_servers_against_team( detail={"error": detail}, ) - # Validate requested toolsets (must be subset of team's toolsets) - if requested_toolsets: + # Validate requested toolsets against team's allowed toolsets. + # Only enforce the team-based restriction when a team is present — standalone + # keys (no team) can freely be granted any toolset by an admin. + if requested_toolsets and team_obj is not None: team_toolsets: Set[str] = set() if ( - team_obj is not None - and team_obj.object_permission is not None + team_obj.object_permission is not None and team_obj.object_permission.mcp_toolsets ): team_toolsets = set(team_obj.object_permission.mcp_toolsets) disallowed_toolsets = requested_toolsets - team_toolsets if disallowed_toolsets: - if team_obj is not None: - detail = ( - f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': " - f"{sorted(disallowed_toolsets)}. " - f"Team allows: {sorted(team_toolsets)}." - ) - else: - detail = ( - f"Key is not in a team. MCP toolsets cannot be assigned to " - f"keys outside of a team. Disallowed toolsets: {sorted(disallowed_toolsets)}." - ) raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail={"error": detail}, + detail={ + "error": ( + f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': " + f"{sorted(disallowed_toolsets)}. " + f"Team allows: {sorted(team_toolsets)}." + ) + }, ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 9712ec7aa55..31b04efb3d7 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -160,28 +160,24 @@ class LiteLLM_Proxy_MCP_Handler: ): mcp_servers.append(server_url.split("/")[-1]) - # Resolve toolset names: if any name in mcp_servers is a toolset (not a real - # server name), apply toolset scope to user_api_key_auth so that only the - # toolset's servers and tools are visible. Non-toolset names are kept as-is. + # Resolve toolset names: collect all toolset IDs first, then apply their + # combined permissions in a single pass so multiple toolsets are unioned + # rather than the last one overwriting the others. resolved_mcp_servers: List[str] = [] + resolved_toolset_ids: List[str] = [] for name in mcp_servers: if not global_mcp_server_manager.get_mcp_server_by_name(name): try: - from litellm.proxy._experimental.mcp_server.server import ( - _apply_toolset_scope, - ) from litellm.proxy.proxy_server import prisma_client - if prisma_client is not None and user_api_key_auth is not None: + if prisma_client is not None: toolset = ( await global_mcp_server_manager.get_toolset_by_name_cached( prisma_client, name ) ) if toolset is not None: - user_api_key_auth = await _apply_toolset_scope( - user_api_key_auth, toolset.toolset_id - ) + resolved_toolset_ids.append(toolset.toolset_id) # Don't add to resolved_mcp_servers — toolset scope # restricts via object_permission, not server name filter. continue @@ -189,6 +185,39 @@ class LiteLLM_Proxy_MCP_Handler: verbose_logger.debug(f"Could not resolve '{name}' as toolset: {_e}") resolved_mcp_servers.append(name) + # Apply all resolved toolsets at once (union), avoiding permission overwrite. + if resolved_toolset_ids and user_api_key_auth is not None: + try: + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + tool_permissions = ( + await global_mcp_server_manager.resolve_toolset_tool_permissions( + toolset_ids=resolved_toolset_ids + ) + ) + server_ids = list(tool_permissions.keys()) + existing_op = user_api_key_auth.object_permission + if existing_op is not None: + updated_op = existing_op.model_copy( + update={ + "mcp_servers": server_ids, + "mcp_tool_permissions": tool_permissions, + "mcp_toolsets": [], + "mcp_access_groups": [], + } + ) + else: + updated_op = LiteLLM_ObjectPermissionTable( + object_permission_id="toolset-scope", + mcp_servers=server_ids, + mcp_tool_permissions=tool_permissions, + ) + user_api_key_auth = user_api_key_auth.model_copy( + update={"object_permission": updated_op} + ) + except Exception as _e: + verbose_logger.debug(f"Could not apply toolset permissions: {_e}") + tools = await _get_tools_from_mcp_servers( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header,