Add hook integration: cluster balancing, subscription routing, tariff menu

handlers/keys/key_mode/key_cluster_mode.py

@@ -92,21 +92,26 @@ async def key_cluster_mode(
                 if tariff.get("traffic_limit") is not None:
                     traffic_limit_gb = int(tariff["traffic_limit"])
 
-        try:
-            least_loaded_cluster = await get_least_loaded_cluster(session)
-        except ValueError as e:
-            logger.error(f"Нет доступных кластеров: {e}")
-            error_message = str(e)
+        forced_cluster_results = await run_hooks("cluster_override", tg_id=tg_id, state_data=data, session=session, plan=plan)
+        
+        if forced_cluster_results and forced_cluster_results[0]:
+            least_loaded_cluster = forced_cluster_results[0]
+        else:
+            try:
+                least_loaded_cluster = await get_least_loaded_cluster(session)
+            except ValueError as e:
+                logger.error(f"Нет доступных кластеров: {e}")
+                error_message = str(e)
 
-            if safe_to_edit:
-                await edit_or_send_message(
-                    target_message=target_message,
-                    text=error_message,
-                    reply_markup=None,
-                )
-            else:
-                await bot.send_message(chat_id=tg_id, text=error_message)
-            return
+                if safe_to_edit:
+                    await edit_or_send_message(
+                        target_message=target_message,
+                        text=error_message,
+                        reply_markup=None,
+                    )
+                else:
+                    await bot.send_message(chat_id=tg_id, text=error_message)
+                return
 
         await create_key_on_cluster(
             cluster_id=least_loaded_cluster,

handlers/keys/key_mode/key_country_mode.py

@@ -78,16 +78,23 @@ async def key_country_mode(
         target_message = message_or_query
         safe_to_edit = True
 
-    try:
-        least_loaded_cluster = await get_least_loaded_cluster(session)
-    except ValueError as e:
-        logger.error(f"Нет доступных кластеров: {e}")
-        text = str(e)
-        if safe_to_edit:
-            await edit_or_send_message(target_message=target_message, text=text, reply_markup=None)
-        else:
-            await bot.send_message(chat_id=tg_id, text=text)
-        return
+    data = await state.get_data() if state else {}
+
+    forced_cluster_results = await run_hooks("cluster_override", tg_id=tg_id, state_data=data, session=session, plan=plan)
+    
+    if forced_cluster_results and forced_cluster_results[0]:
+        least_loaded_cluster = forced_cluster_results[0]
+    else:
+        try:
+            least_loaded_cluster = await get_least_loaded_cluster(session)
+        except ValueError as e:
+            logger.error(f"Нет доступных кластеров: {e}")
+            text = str(e)
+            if safe_to_edit:
+                await edit_or_send_message(target_message=target_message, text=text, reply_markup=None)
+            else:
+                await bot.send_message(chat_id=tg_id, text=text)
+            return
 
     result = await session.execute(
         select(

handlers/keys/key_mode/key_create.py

@@ -39,6 +39,8 @@ from handlers.texts import (
     SELECT_TARIFF_PLAN_MSG,
 )
 from handlers.utils import edit_or_send_message, format_discount_time_left, get_least_loaded_cluster
+from hooks.hook_buttons import insert_hook_buttons
+from hooks.hooks import run_hooks
 from logger import logger
 
 from .key_cluster_mode import key_cluster_mode
@@ -240,6 +242,9 @@ async def handle_key_creation(
                 )
             )
 
+        tariff_menu_buttons = await run_hooks("tariff_menu", group_code=group_code, cluster_name=cluster_name, tg_id=tg_id, session=session)
+        builder = insert_hook_buttons(builder, tariff_menu_buttons)
+
         builder.row(InlineKeyboardButton(text=MAIN_MENU, callback_data="profile"))
 
         target_message = message_or_query.message if isinstance(message_or_query, CallbackQuery) else message_or_query

handlers/utils.py

@@ -21,6 +21,7 @@ from bot import bot
 from config import ADMIN_ID
 from database import get_servers
 from database.models import Key, Notification, Server
+from hooks.hooks import run_hooks
 from logger import logger
 
 
@@ -67,6 +68,10 @@ async def get_least_loaded_cluster(session: AsyncSession) -> str:
         else:
             continue
 
+    cluster_filter_results = await run_hooks("cluster_balancer", available_clusters=available_clusters, session=session)
+    if cluster_filter_results and cluster_filter_results[0]:
+        available_clusters = cluster_filter_results[0]
+
     if not available_clusters:
         logger.warning("❌ Нет доступных кластеров с лимитом ключей!")
         raise ValueError("⚠️ Сервисы временно недоступны. Попробуйте позже.")