fix: 增加请求前限额拦截(async_pre_call_hook)，修复超限滞后击穿漏洞

volc_limit_sync.py

@@ -17,7 +17,11 @@ from urllib.parse import urlparse
 try:
     import litellm
     from litellm.integrations.custom_logger import CustomLogger
-    from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError
+    from litellm.exceptions import (
+        ServiceUnavailableError,
+        ContextWindowExceededError,
+        RateLimitError,
+    )
 
     LITELLM_AVAILABLE = True
 except ImportError:
@@ -33,6 +37,9 @@ except ImportError:
     class ContextWindowExceededError(Exception):
         pass
 
+    class RateLimitError(Exception):
+        pass
+
 
 # ==========================================
 # 🛠️ Configuration & Constants
@@ -626,6 +633,55 @@ def startup_audit():
         log(f"❌ Startup audit failed: {e}", level="ERROR")
 
 
+def _check_request_block(model_name):
+    """
+    请求前闸门：对外模型名 → 判定是否超限。
+    只拦截 MONITORED_MODEL_PAIRS 内的火山监控条目；newapi 等非监控渠道天然放行。
+    返回 (clean_model, usage, limit) 表示超限，None 表示放行。任何异常 fail-open。
+    """
+    if not redis_client or not LITELLM_AVAILABLE:
+        return None
+    try:
+        from litellm.proxy.proxy_server import llm_router
+
+        if not llm_router or not llm_router.model_list:
+            return None
+
+        for entry in llm_router.model_list:
+            if entry.get("model_name") != model_name:
+                continue
+            # 已被限额踢出的条目跳过（路由层已匹配不到）
+            if "_daily_limit_backup" in entry:
+                continue
+
+            params = entry.get("litellm_params", {})
+            raw_model = params.get("model", "")
+            api_base = params.get("api_base", "")
+
+            clean = clean_model_name(raw_model)
+            h = generate_api_base_hash(api_base)
+
+            # 只拦已纳入监控的火山模型（模型 + 端点粒度）
+            if MONITORED_MODEL_PAIRS and (clean, h) not in MONITORED_MODEL_PAIRS:
+                continue
+
+            effective_limit = get_effective_limit(clean, api_base)
+            if effective_limit is None:
+                continue
+
+            usage_key = get_usage_key(clean, api_base)
+            val = redis_client.get(usage_key)
+            usage = float(val or 0)
+
+            if usage >= float(effective_limit):
+                return (clean, usage, effective_limit)
+
+        return None
+    except Exception as e:
+        log(f"❌ Pre-block check failed: {e}", level="ERROR")
+        return None
+
+
 def quota_interceptor(kwargs, user_api_key_dict=None):
     try:
         # 1. 尝试恢复（定期）
@@ -646,6 +702,32 @@ class DailyLimitHandler(CustomLogger):
                 litellm.callbacks.append(self)
                 dprint("✅ DailyLimitHandler registered")
 
+    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+        """
+        请求前限额闸门：在 route_request 之前拒绝超限请求（429），
+        不进入 router → 不触发 num_retries / fallback 放大。
+        """
+        if call_type not in ("completion", "acompletion", "text_completion"):
+            return None
+        if not data or not data.get("model"):
+            return None
+
+        blocked = _check_request_block(data["model"])
+        if not blocked:
+            return None
+
+        clean, usage, limit = blocked
+        log(
+            f"🚫 [PRE-BLOCK] {data['model']} daily limit reached ({int(usage):,}/{int(limit):,})",
+            level="WARNING",
+        )
+        raise RateLimitError(
+            message=f"Daily limit reached for model '{data['model']}' ({int(usage):,}/{int(limit):,})",
+            model=data["model"],
+            llm_provider="volcengine",
+            request_data=data,
+        )
+
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         self._record_usage(kwargs, response_obj)