File size: 6,570 Bytes
b461de0 5039d41 b461de0 5039d41 b461de0 5039d41 b461de0 5039d41 b461de0 5039d41 b461de0 5039d41 b461de0 178e3be b461de0 178e3be b461de0 5039d41 b461de0 5039d41 b461de0 178e3be b461de0 5039d41 b461de0 5039d41 b461de0 5039d41 b461de0 5039d41 b461de0 178e3be b461de0 5039d41 b461de0 178e3be b461de0 5039d41 b461de0 178e3be b461de0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | 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()
|