import time import asyncio from typing import Optional, Dict, Any from core.config import ConfigManager from core.noke_client import NokeClient COOLDOWN_SECONDS = 300 MAX_CONSECUTIVE_ERRORS = 3 # 按错误类别冷却时间(秒) COOLDOWN_BY_CATEGORY = { "auth": 3600, # 认证失败,长冷却 "rate_limit": 60, # 限流,短冷却 "signature": 60, # 签名错误,短冷却 "transient": 300, # 瞬态错误,默认 } class TokenManager: def __init__(self, config: ConfigManager): self.config = config self._clients: dict[str, NokeClient] = {} self._index: int = 0 self._cooldowns: dict[str, float] = {} # token_id -> 冷却截止时间戳 self._error_categories: dict[str, str] = {} # token_id -> 最近错误类别 self._lock = asyncio.Lock() self._dirty = False @property def has_tokens(self) -> bool: return len(self.config.get("tokens", default=[])) > 0 def _get_available_tokens(self) -> list[dict]: tokens = self.config.get("tokens", default=[]) now = time.time() available = [] for t in tokens: if not t.get("enabled", True): continue cooldown_until = self._cooldowns.get(t["id"], 0) if now >= cooldown_until: if t["id"] in self._cooldowns: del self._cooldowns[t["id"]] self._error_categories.pop(t["id"], None) t["status"] = "unknown" t["error_count"] = 0 available.append(t) return available def _get_client(self, token_info: dict) -> NokeClient: tid = token_info["id"] if tid not in self._clients: verify_tls = bool(self.config.get("noke", "verify_tls", default=True)) self._clients[tid] = NokeClient( token_info["value"], self.config.get("noke", "base_url"), self.config.get("noke", "client_id"), verify_tls=verify_tls, ) return self._clients[tid] async def get_next(self) -> tuple[Optional[dict], Optional[NokeClient]]: async with self._lock: available = self._get_available_tokens() if not available: return None, None self._index = self._index % len(available) token_info = available[self._index] self._index = (self._index + 1) % len(available) client = self._get_client(token_info) return token_info, client async def get_by_id(self, token_id: str) -> tuple[Optional[dict], Optional[NokeClient]]: """按 token_id 获取特定 token 和 client(用于会话粘性)""" async with self._lock: tokens = self.config.get("tokens", default=[]) for t in tokens: if t["id"] == token_id: # 检查是否可用(未冷却) cooldown_until = self._cooldowns.get(token_id, 0) if time.time() < cooldown_until: return None, None client = self._get_client(t) return t, client return None, None def report_success(self, token_id: str): for t in self.config.get("tokens", default=[]): if t["id"] == token_id: t["total_requests"] = t.get("total_requests", 0) + 1 t["error_count"] = 0 t["last_used_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) t["status"] = "active" break self._dirty = True def report_error(self, token_id: str, category: str = "transient"): """上报错误,按类别决定冷却策略。 - auth/rate_limit/signature: 立即进入冷却(冷却时间按类别) - transient: 累计错误次数,达到 MAX_CONSECUTIVE_ERRORS 才冷却(COOLDOWN_SECONDS) """ for t in self.config.get("tokens", default=[]): if t["id"] == token_id: t["error_count"] = t.get("error_count", 0) + 1 t["total_requests"] = t.get("total_requests", 0) + 1 t["status"] = "error" self._error_categories[token_id] = category # 立即冷却类别 if category in ("auth", "rate_limit", "signature"): cooldown = COOLDOWN_BY_CATEGORY.get(category, COOLDOWN_SECONDS) self._cooldowns[token_id] = time.time() + cooldown t["status"] = "cooldown" t["error_count"] = 0 # 重置计数 else: # transient: 累计到阈值才冷却 if t["error_count"] >= MAX_CONSECUTIVE_ERRORS: self._cooldowns[token_id] = time.time() + COOLDOWN_SECONDS t["status"] = "cooldown" t["error_count"] = 0 break self._dirty = True def remove_client(self, token_id: str): self._clients.pop(token_id, None) self._cooldowns.pop(token_id, None) self._error_categories.pop(token_id, None) def flush(self): if self._dirty: self.config.save() self._dirty = False def get_token_status(self, token_id: str) -> str: now = time.time() cooldown_until = self._cooldowns.get(token_id, 0) if now < cooldown_until: return "cooldown" for t in self.config.get("tokens", default=[]): if t["id"] == token_id: return t.get("status", "unknown") return "unknown" def snapshot(self) -> dict: """返回 token 池快照:总数、可用数、冷却中数量""" tokens = self.config.get("tokens", default=[]) total = len(tokens) now = time.time() available = 0 cooldown = 0 for t in tokens: if not t.get("enabled", True): continue cooldown_until = self._cooldowns.get(t["id"], 0) if now < cooldown_until: cooldown += 1 else: available += 1 return { "total": total, "available": available, "cooldown": cooldown, } async def close_all(self): self.flush() for client in self._clients.values(): await client.client.aclose() self._clients.clear()