"""Grok Token 管理器 - 单例模式的Token负载均衡和状态管理""" import orjson import time import asyncio import aiofiles import portalocker from pathlib import Path from curl_cffi.requests import AsyncSession from typing import Dict, Any, Optional, Tuple from app.models.grok_models import TokenType, Models from app.core.exception import GrokApiException from app.core.logger import logger from app.core.config import setting from app.services.grok.statsig import get_dynamic_headers # 常量 RATE_LIMIT_API = "https://grok.com/rest/rate-limits" TIMEOUT = 30 BROWSER = "chrome133a" MAX_FAILURES = 3 TOKEN_INVALID = 401 STATSIG_INVALID = 403 # 冷却常量 COOLDOWN_REQUESTS = 5 # 普通失败冷却请求数 COOLDOWN_429_WITH_QUOTA = 3600 # 429+有额度冷却1小时（秒） COOLDOWN_429_NO_QUOTA = 36000 # 429+无额度冷却10小时（秒） class GrokTokenManager: """Token管理器（单例）""" _instance: Optional['GrokTokenManager'] = None _lock = asyncio.Lock() def __new__(cls) -> 'GrokTokenManager': if cls._instance is None: cls._instance = super().__new__(cls) return cls._instance def __init__(self): if hasattr(self, '_initialized'): return self.token_file = Path(__file__).parents[3] / "data" / "token.json" self._file_lock = asyncio.Lock() self.token_file.parent.mkdir(parents=True, exist_ok=True) self._storage = None self.token_data = None # 延迟加载 # 批量保存队列 self._save_pending = False # 标记是否有待保存的数据 self._save_task = None # 后台保存任务 self._shutdown = False # 关闭标志 # 冷却状态 self._cooldown_counts: Dict[str, int] = {} # Token -> 剩余冷却次数 self._request_counter = 0 # 全局请求计数器 # 刷新状态 self._refresh_lock = False # 刷新锁 self._refresh_progress: Dict[str, Any] = {"running": False, "current": 0, "total": 0, "success": 0, "failed": 0} self._initialized = True logger.debug(f"[Token] 初始化完成: {self.token_file}") def set_storage(self, storage) -> None: """设置存储实例""" self._storage = storage async def _load_data(self) -> None: """异步加载Token数据（支持多进程）""" default = {TokenType.NORMAL.value: {}, TokenType.SUPER.value: {}} def load_sync(): with open(self.token_file, "r", encoding="utf-8") as f: portalocker.lock(f, portalocker.LOCK_SH) try: return orjson.loads(f.read()) finally: portalocker.unlock(f) try: if self.token_file.exists(): # 使用进程锁读取文件 async with self._file_lock: self.token_data = await asyncio.to_thread(load_sync) else: self.token_data = default logger.debug("[Token] 创建新数据文件") except Exception as e: logger.error(f"[Token] 加载失败: {e}") self.token_data = default async def _save_data(self) -> None: """保存Token数据（支持多进程）""" def save_sync(data): with open(self.token_file, "w", encoding="utf-8") as f: portalocker.lock(f, portalocker.LOCK_EX) try: content = orjson.dumps(data, option=orjson.OPT_INDENT_2).decode() f.write(content) f.flush() finally: portalocker.unlock(f) try: if not self._storage: async with self._file_lock: await asyncio.to_thread(save_sync, self.token_data) else: await self._storage.save_tokens(self.token_data) except Exception as e: logger.error(f"[Token] 保存失败: {e}") raise GrokApiException(f"保存失败: {e}", "TOKEN_SAVE_ERROR") def _mark_dirty(self) -> None: """标记有待保存的数据""" self._save_pending = True async def _batch_save_worker(self) -> None: """批量保存后台任务""" from app.core.config import setting interval = setting.global_config.get("batch_save_interval", 1.0) logger.info(f"[Token] 存储任务已启动，间隔: {interval}s") while not self._shutdown: await asyncio.sleep(interval) if self._save_pending and not self._shutdown: try: await self._save_data() self._save_pending = False logger.debug("[Token] 存储完成") except Exception as e: logger.error(f"[Token] 存储失败: {e}") async def start_batch_save(self) -> None: """启动批量保存任务""" if self._save_task is None: self._save_task = asyncio.create_task(self._batch_save_worker()) logger.info("[Token] 存储任务已创建") async def shutdown(self) -> None: """关闭并刷新所有待保存数据""" self._shutdown = True if self._save_task: self._save_task.cancel() try: await self._save_task except asyncio.CancelledError: pass # 最终刷新 if self._save_pending: await self._save_data() logger.info("[Token] 关闭时刷新完成") @staticmethod def _extract_sso(auth_token: str) -> Optional[str]: """提取SSO值""" if "sso=" in auth_token: return auth_token.split("sso=")[1].split(";")[0] logger.warning("[Token] 无法提取SSO值") return None def _find_token(self, sso: str) -> Tuple[Optional[str], Optional[Dict]]: """查找Token""" for token_type in [TokenType.NORMAL.value, TokenType.SUPER.value]: if sso in self.token_data[token_type]: return token_type, self.token_data[token_type][sso] return None, None async def add_token(self, tokens: list[str], token_type: TokenType) -> None: """添加Token""" if not tokens: return count = 0 for token in tokens: if not token or not token.strip(): continue self.token_data[token_type.value][token] = { "createdTime": int(time.time() * 1000), "remainingQueries": -1, "heavyremainingQueries": -1, "status": "active", "failedCount": 0, "lastFailureTime": None, "lastFailureReason": None, "tags": [], "note": "" } count += 1 self._mark_dirty() # 批量保存 logger.info(f"[Token] 添加 {count} 个 {token_type.value} Token") async def delete_token(self, tokens: list[str], token_type: TokenType) -> None: """删除Token""" if not tokens: return count = 0 for token in tokens: if token in self.token_data[token_type.value]: del self.token_data[token_type.value][token] count += 1 self._mark_dirty() # 批量保存 logger.info(f"[Token] 删除 {count} 个 {token_type.value} Token") async def update_token_tags(self, token: str, token_type: TokenType, tags: list[str]) -> None: """更新Token标签""" if token not in self.token_data[token_type.value]: raise GrokApiException("Token不存在", "TOKEN_NOT_FOUND", {"token": token[:10]}) cleaned = [t.strip() for t in tags if t and t.strip()] self.token_data[token_type.value][token]["tags"] = cleaned self._mark_dirty() # 批量保存 logger.info(f"[Token] 更新标签: {token[:10]}... -> {cleaned}") async def update_token_note(self, token: str, token_type: TokenType, note: str) -> None: """更新Token备注""" if token not in self.token_data[token_type.value]: raise GrokApiException("Token不存在", "TOKEN_NOT_FOUND", {"token": token[:10]}) self.token_data[token_type.value][token]["note"] = note.strip() self._mark_dirty() # 批量保存 logger.info(f"[Token] 更新备注: {token[:10]}...") def get_tokens(self) -> Dict[str, Any]: """获取所有Token""" return self.token_data.copy() async def _reload_if_needed(self) -> None: """在多进程模式下重新加载数据""" # 只在文件模式且多进程环境下才重新加载 if self._storage: return def reload_sync(): with open(self.token_file, "r", encoding="utf-8") as f: portalocker.lock(f, portalocker.LOCK_SH) try: return orjson.loads(f.read()) finally: portalocker.unlock(f) try: if self.token_file.exists(): self.token_data = await asyncio.to_thread(reload_sync) except Exception as e: logger.warning(f"[Token] 重新加载失败: {e}") async def get_token(self, model: str) -> str: """获取Token""" jwt = await self.select_token(model) return f"sso-rw={jwt};sso={jwt}" async def select_token(self, model: str) -> str: """选择最优Token（多进程安全，支持冷却）""" # 重新加载最新数据（多进程模式） await self._reload_if_needed() # 递减所有次数冷却计数 self._request_counter += 1 for token in list(self._cooldown_counts.keys()): self._cooldown_counts[token] -= 1 if self._cooldown_counts[token] <= 0: del self._cooldown_counts[token] logger.debug(f"[Token] 冷却结束: {token[:10]}...") current_time = time.time() * 1000 # 毫秒 def select_best(tokens: Dict[str, Any], field: str) -> Tuple[Optional[str], Optional[int]]: """选择最佳Token""" unused, used = [], [] for key, data in tokens.items(): # 跳过已失效的token if data.get("status") == "expired": continue # 跳过失败次数过多的token（任何错误状态码） if data.get("failedCount", 0) >= MAX_FAILURES: continue # 跳过次数冷却中的token if key in self._cooldown_counts: continue # 跳过时间冷却中的token（429） cooldown_until = data.get("cooldownUntil", 0) if cooldown_until and cooldown_until > current_time: continue remaining = int(data.get(field, -1)) if remaining == 0: continue if remaining == -1: unused.append(key) elif remaining > 0: used.append((key, remaining)) if unused: return unused[0], -1 if used: used.sort(key=lambda x: x[1], reverse=True) return used[0][0], used[0][1] return None, None # 快照 snapshot = { TokenType.NORMAL.value: self.token_data[TokenType.NORMAL.value].copy(), TokenType.SUPER.value: self.token_data[TokenType.SUPER.value].copy() } # 选择策略 if model == "grok-4-heavy": field = "heavyremainingQueries" token_key, remaining = select_best(snapshot[TokenType.SUPER.value], field) else: field = "remainingQueries" token_key, remaining = select_best(snapshot[TokenType.NORMAL.value], field) if token_key is None: token_key, remaining = select_best(snapshot[TokenType.SUPER.value], field) if token_key is None: raise GrokApiException( f"没有可用Token: {model}", "NO_AVAILABLE_TOKEN", { "model": model, "normal": len(snapshot[TokenType.NORMAL.value]), "super": len(snapshot[TokenType.SUPER.value]), "cooldown_count": len(self._cooldown_counts) } ) status = "未使用" if remaining == -1 else f"剩余{remaining}次" logger.debug(f"[Token] 分配Token: {model} ({status})") return token_key async def check_limits(self, auth_token: str, model: str) -> Optional[Dict[str, Any]]: """检查速率限制""" try: rate_model = Models.to_rate_limit(model) payload = {"requestKind": "DEFAULT", "modelName": rate_model} cf = setting.grok_config.get("cf_clearance", "") headers = get_dynamic_headers("/rest/rate-limits") headers["Cookie"] = f"{auth_token};{cf}" if cf else auth_token # 外层重试：可配置状态码（401/429等） retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) MAX_OUTER_RETRY = 3 for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次 # 内层重试：403代理池重试 max_403_retries = 5 retry_403_count = 0 while retry_403_count <= max_403_retries: # 异步获取代理（支持代理池） from app.core.proxy_pool import proxy_pool # 如果是403重试且使用代理池，强制刷新代理 if retry_403_count > 0 and proxy_pool._enabled: logger.info(f"[Token] 403重试 {retry_403_count}/{max_403_retries}，刷新代理...") proxy = await proxy_pool.force_refresh() else: proxy = await setting.get_proxy_async("service") proxies = {"http": proxy, "https": proxy} if proxy else None async with AsyncSession() as session: response = await session.post( RATE_LIMIT_API, headers=headers, json=payload, impersonate=BROWSER, timeout=TIMEOUT, proxies=proxies ) # 内层403重试：仅当有代理池时触发 if response.status_code == 403 and proxy_pool._enabled: retry_403_count += 1 if retry_403_count <= max_403_retries: logger.warning(f"[Token] 遇到403错误，正在重试 ({retry_403_count}/{max_403_retries})...") await asyncio.sleep(0.5) continue # 内层重试全部失败 logger.error(f"[Token] 403错误，已重试{retry_403_count-1}次，放弃") sso = self._extract_sso(auth_token) if sso: await self.record_failure(auth_token, 403, "服务器被Block") # 检查可配置状态码错误 - 外层重试 if response.status_code in retry_codes: if outer_retry < MAX_OUTER_RETRY: delay = (outer_retry + 1) * 0.1 # 渐进延迟：0.1s, 0.2s, 0.3s logger.warning(f"[Token] 遇到{response.status_code}错误，外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})，等待{delay}s...") await asyncio.sleep(delay) break # 跳出内层循环，进入外层重试 else: logger.error(f"[Token] {response.status_code}错误，已重试{outer_retry}次，放弃") sso = self._extract_sso(auth_token) if sso: if response.status_code == 401: await self.record_failure(auth_token, 401, "Token失效") else: await self.record_failure(auth_token, response.status_code, f"错误: {response.status_code}") return None if response.status_code == 200: data = response.json() sso = self._extract_sso(auth_token) if outer_retry > 0 or retry_403_count > 0: logger.info(f"[Token] 重试成功！") if sso: if model == "grok-4-heavy": await self.update_limits(sso, normal=None, heavy=data.get("remainingQueries", -1)) logger.info(f"[Token] 更新限制: {sso[:10]}..., heavy={data.get('remainingQueries', -1)}") else: await self.update_limits(sso, normal=data.get("remainingTokens", -1), heavy=None) logger.info(f"[Token] 更新限制: {sso[:10]}..., basic={data.get('remainingTokens', -1)}") return data else: # 其他错误 logger.warning(f"[Token] 获取限制失败: {response.status_code}") sso = self._extract_sso(auth_token) if sso: await self.record_failure(auth_token, response.status_code, f"错误: {response.status_code}") return None except Exception as e: logger.error(f"[Token] 检查限制错误: {e}") return None async def update_limits(self, sso: str, normal: Optional[int] = None, heavy: Optional[int] = None) -> None: """更新限制""" try: for token_type in [TokenType.NORMAL.value, TokenType.SUPER.value]: if sso in self.token_data[token_type]: if normal is not None: self.token_data[token_type][sso]["remainingQueries"] = normal if heavy is not None: self.token_data[token_type][sso]["heavyremainingQueries"] = heavy self._mark_dirty() # 批量保存 logger.info(f"[Token] 更新限制: {sso[:10]}...") return logger.warning(f"[Token] 未找到: {sso[:10]}...") except Exception as e: logger.error(f"[Token] 更新限制错误: {e}") async def record_failure(self, auth_token: str, status: int, msg: str) -> None: """记录失败""" try: if status == STATSIG_INVALID: logger.warning("[Token] IP被Block，请: 1.更换IP 2.使用代理 3.配置CF值") return sso = self._extract_sso(auth_token) if not sso: return _, data = self._find_token(sso) if not data: logger.warning(f"[Token] 未找到: {sso[:10]}...") return data["failedCount"] = data.get("failedCount", 0) + 1 data["lastFailureTime"] = int(time.time() * 1000) data["lastFailureReason"] = f"{status}: {msg}" logger.warning( f"[Token] 失败: {sso[:10]}... (状态:{status}), " f"次数: {data['failedCount']}/{MAX_FAILURES}, 原因: {msg}" ) if 400 <= status < 500 and data["failedCount"] >= MAX_FAILURES: data["status"] = "expired" logger.error(f"[Token] 标记失效: {sso[:10]}... (连续{status}错误{data['failedCount']}次)") self._mark_dirty() # 批量保存 except Exception as e: logger.error(f"[Token] 记录失败错误: {e}") async def reset_failure(self, auth_token: str) -> None: """重置失败计数""" try: sso = self._extract_sso(auth_token) if not sso: return _, data = self._find_token(sso) if not data: return if data.get("failedCount", 0) > 0: data["failedCount"] = 0 data["lastFailureTime"] = None data["lastFailureReason"] = None self._mark_dirty() # 批量保存 logger.info(f"[Token] 重置失败计数: {sso[:10]}...") except Exception as e: logger.error(f"[Token] 重置失败错误: {e}") async def apply_cooldown(self, auth_token: str, status_code: int) -> None: """应用冷却策略 - 429 错误：使用时间冷却（有额度1小时，无额度10小时） - 其他错误：使用次数冷却（5次请求） """ try: sso = self._extract_sso(auth_token) if not sso: return _, data = self._find_token(sso) if not data: return remaining = data.get("remainingQueries", -1) if status_code == 429: # 429 使用时间冷却 if remaining > 0 or remaining == -1: # 有额度：冷却1小时 cooldown_until = time.time() + COOLDOWN_429_WITH_QUOTA logger.info(f"[Token] 429冷却(有额度): {sso[:10]}... 冷却1小时") else: # 无额度：冷却10小时 cooldown_until = time.time() + COOLDOWN_429_NO_QUOTA logger.info(f"[Token] 429冷却(无额度): {sso[:10]}... 冷却10小时") data["cooldownUntil"] = int(cooldown_until * 1000) self._mark_dirty() else: # 其他错误使用次数冷却（有额度时才冷却） if remaining != 0: self._cooldown_counts[sso] = COOLDOWN_REQUESTS logger.info(f"[Token] 次数冷却: {sso[:10]}... 冷却{COOLDOWN_REQUESTS}次请求") except Exception as e: logger.error(f"[Token] 应用冷却错误: {e}") async def refresh_all_limits(self) -> Dict[str, Any]: """刷新所有 Token 的剩余次数""" # 检查是否已在刷新 if self._refresh_lock: return {"error": "refresh_in_progress", "message": "已有刷新任务在进行中", "progress": self._refresh_progress} # 获取锁 self._refresh_lock = True try: # 计算总数 all_tokens = [] for token_type in [TokenType.NORMAL.value, TokenType.SUPER.value]: for sso in list(self.token_data[token_type].keys()): all_tokens.append((token_type, sso)) total = len(all_tokens) self._refresh_progress = {"running": True, "current": 0, "total": total, "success": 0, "failed": 0} success_count = 0 fail_count = 0 for i, (token_type, sso) in enumerate(all_tokens): auth_token = f"sso-rw={sso};sso={sso}" try: result = await self.check_limits(auth_token, "grok-4-fast") if result: success_count += 1 else: fail_count += 1 except Exception as e: logger.warning(f"[Token] 刷新失败: {sso[:10]}... - {e}") fail_count += 1 # 更新进度 self._refresh_progress = { "running": True, "current": i + 1, "total": total, "success": success_count, "failed": fail_count } await asyncio.sleep(0.1) # 避免请求过快 logger.info(f"[Token] 批量刷新完成: 成功{success_count}, 失败{fail_count}") self._refresh_progress = {"running": False, "current": total, "total": total, "success": success_count, "failed": fail_count} return {"success": success_count, "failed": fail_count, "total": total} finally: self._refresh_lock = False def get_refresh_progress(self) -> Dict[str, Any]: """获取刷新进度""" return self._refresh_progress.copy() # 全局实例 token_manager = GrokTokenManager()