"""Token 管理服务""" import asyncio import time from datetime import datetime from typing import Dict, List, Optional from app.core.logger import logger from app.services.token.models import ( TokenInfo, EffortType, FAIL_THRESHOLD, TokenStatus, BASIC__DEFAULT_QUOTA, SUPER_DEFAULT_QUOTA, ) from app.core.storage import get_storage, LocalStorage from app.core.config import get_config from app.core.exceptions import UpstreamException from app.services.token.pool import TokenPool from app.services.grok.batch_services.usage import UsageService DEFAULT_REFRESH_BATCH_SIZE = 10 DEFAULT_REFRESH_CONCURRENCY = 5 DEFAULT_SUPER_REFRESH_INTERVAL_HOURS = 2 DEFAULT_REFRESH_INTERVAL_HOURS = 8 DEFAULT_RELOAD_INTERVAL_SEC = 30 DEFAULT_SAVE_DELAY_MS = 500 DEFAULT_USAGE_FLUSH_INTERVAL_SEC = 5 SUPER_POOL_NAME = "ssoSuper" BASIC_POOL_NAME = "ssoBasic" def _default_quota_for_pool(pool_name: str) -> int: if pool_name == SUPER_POOL_NAME: return SUPER_DEFAULT_QUOTA return BASIC__DEFAULT_QUOTA class TokenManager: """管理 Token 的增删改查和配额同步""" _instance: Optional["TokenManager"] = None _lock = asyncio.Lock() def __init__(self): self.pools: Dict[str, TokenPool] = {} self.initialized = False self._save_lock = asyncio.Lock() self._dirty = False self._save_task: Optional[asyncio.Task] = None self._save_delay = DEFAULT_SAVE_DELAY_MS / 1000.0 self._last_reload_at = 0.0 self._has_state_changes = False self._has_usage_changes = False self._state_change_seq = 0 self._usage_change_seq = 0 self._last_usage_flush_at = 0.0 self._dirty_tokens = {} self._dirty_deletes = set() @classmethod async def get_instance(cls) -> "TokenManager": """获取单例实例""" if cls._instance is None: async with cls._lock: if cls._instance is None: cls._instance = cls() await cls._instance._load() return cls._instance async def _load(self): """初始化加载""" if not self.initialized: try: storage = get_storage() data = await storage.load_tokens() # 如果后端返回 None 或空数据，尝试从本地 data/token.json 初始化后端 if not data: local_storage = LocalStorage() local_data = await local_storage.load_tokens() if local_data: data = local_data await storage.save_tokens(local_data) logger.info( f"Initialized remote token storage ({storage.__class__.__name__}) with local tokens." ) else: data = {} self.pools = {} for pool_name, tokens in data.items(): pool = TokenPool(pool_name) for token_data in tokens: quota_missing = not ( isinstance(token_data, dict) and "quota" in token_data ) try: # 统一存储裸 token if isinstance(token_data, dict): raw_token = token_data.get("token") if isinstance(raw_token, str) and raw_token.startswith( "sso=" ): token_data["token"] = raw_token[4:] token_info = TokenInfo(**token_data) if quota_missing and pool_name == SUPER_POOL_NAME: token_info.quota = SUPER_DEFAULT_QUOTA pool.add(token_info) except Exception as e: logger.warning( f"Failed to load token in pool '{pool_name}': {e}" ) continue pool._rebuild_index() self.pools[pool_name] = pool self.initialized = True self._last_reload_at = time.monotonic() total = sum(p.count() for p in self.pools.values()) logger.info( f"TokenManager initialized: {len(self.pools)} pools with {total} tokens" ) except Exception as e: logger.error(f"Failed to initialize TokenManager: {e}") self.pools = {} self.initialized = True async def reload(self): """重新加载 Token 池数据""" async with self.__class__._lock: self.initialized = False await self._load() async def reload_if_stale(self): """在多 worker 场景下保持短周期一致性""" interval = get_config("token.reload_interval_sec", DEFAULT_RELOAD_INTERVAL_SEC) try: interval = float(interval) except Exception: interval = float(DEFAULT_RELOAD_INTERVAL_SEC) if interval <= 0: return if time.monotonic() - self._last_reload_at < interval: return await self.reload() def _mark_state_change(self): self._has_state_changes = True self._state_change_seq += 1 def _mark_usage_change(self): self._has_usage_changes = True self._usage_change_seq += 1 def _track_token_change( self, token: TokenInfo, pool_name: str, change_kind: str ): token_key = token.token if token_key.startswith("sso="): token_key = token_key[4:] if token_key in self._dirty_deletes: self._dirty_deletes.remove(token_key) existing = self._dirty_tokens.get(token_key) if existing and existing[1] == "state": change_kind = "state" self._dirty_tokens[token_key] = (pool_name, change_kind) if change_kind == "state": self._mark_state_change() else: self._mark_usage_change() def _track_token_delete(self, token_str: str): token_key = token_str if token_key.startswith("sso="): token_key = token_key[4:] self._dirty_deletes.add(token_key) if token_key in self._dirty_tokens: del self._dirty_tokens[token_key] self._mark_state_change() async def _save(self, force: bool = False): """保存变更""" async with self._save_lock: try: if not self._dirty_tokens and not self._dirty_deletes: return if not force and not self._has_state_changes: interval_sec = get_config( "token.usage_flush_interval_sec", DEFAULT_USAGE_FLUSH_INTERVAL_SEC, ) try: interval_sec = float(interval_sec) except Exception: interval_sec = float(DEFAULT_USAGE_FLUSH_INTERVAL_SEC) if interval_sec > 0: now = time.monotonic() if now - self._last_usage_flush_at < interval_sec: self._dirty = True return state_seq = self._state_change_seq usage_seq = self._usage_change_seq dirty_tokens = self._dirty_tokens dirty_deletes = self._dirty_deletes self._dirty_tokens = {} self._dirty_deletes = set() updates = [] deleted = list(dirty_deletes) for token_key, meta in dirty_tokens.items(): if token_key in dirty_deletes: continue pool_name, change_kind = meta pool = self.pools.get(pool_name) if not pool: continue info = pool.get(token_key) if not info: continue payload = info.model_dump() payload["pool_name"] = pool_name payload["_update_kind"] = change_kind updates.append(payload) storage = get_storage() async with storage.acquire_lock("tokens_save", timeout=10): await storage.save_tokens_delta(updates, deleted) if state_seq == self._state_change_seq: self._has_state_changes = False if usage_seq == self._usage_change_seq: self._has_usage_changes = False self._last_usage_flush_at = time.monotonic() except Exception as e: logger.error(f"Failed to save tokens: {e}") self._dirty = True if 'dirty_tokens' in locals(): for token_key, meta in dirty_tokens.items(): existing = self._dirty_tokens.get(token_key) if existing and existing[1] == "state": continue if meta[1] == "state" and existing: self._dirty_tokens[token_key] = (meta[0], "state") else: self._dirty_tokens[token_key] = meta self._dirty_deletes.update(dirty_deletes) for token_key in dirty_deletes: if token_key in self._dirty_tokens: del self._dirty_tokens[token_key] def _schedule_save(self): """合并高频保存请求，减少写入开销""" delay_ms = get_config("token.save_delay_ms", DEFAULT_SAVE_DELAY_MS) try: delay_ms = float(delay_ms) except Exception: delay_ms = float(DEFAULT_SAVE_DELAY_MS) self._save_delay = max(0.0, delay_ms / 1000.0) self._dirty = True if self._save_delay == 0: if self._save_task and not self._save_task.done(): return self._save_task = asyncio.create_task(self._save()) return if self._save_task and not self._save_task.done(): return self._save_task = asyncio.create_task(self._flush_loop()) async def _flush_loop(self): try: while True: await asyncio.sleep(self._save_delay) if not self._dirty: break self._dirty = False await self._save() finally: self._save_task = None if self._dirty: self._schedule_save() def get_token(self, pool_name: str = "ssoBasic", exclude: set = None) -> Optional[str]: """ 获取可用 Token Args: pool_name: Token 池名称 exclude: 需要排除的 token 字符串集合 Returns: Token 字符串或 None """ pool = self.pools.get(pool_name) if not pool: logger.warning(f"Pool '{pool_name}' not found") return None token_info = pool.select(exclude=exclude) if not token_info: logger.warning(f"No available token in pool '{pool_name}'") return None token = token_info.token if token.startswith("sso="): return token[4:] return token def get_token_info(self, pool_name: str = "ssoBasic") -> Optional["TokenInfo"]: """ 获取可用 Token 的完整信息 Args: pool_name: Token 池名称 Returns: TokenInfo 对象或 None """ pool = self.pools.get(pool_name) if not pool: logger.warning(f"Pool '{pool_name}' not found") return None token_info = pool.select() if not token_info: logger.warning(f"No available token in pool '{pool_name}'") return None return token_info def get_token_for_video( self, resolution: str = "480p", video_length: int = 6, pool_candidates: Optional[List[str]] = None, ) -> Optional["TokenInfo"]: """ 根据视频需求智能选择 Token 池 路由策略: - 如果 resolution 是 "720p" 或 video_length > 6: 优先使用 "ssoSuper" 池 - 否则优先使用 "ssoBasic" 池 - 当提供 pool_candidates 时，按候选池顺序回退 Args: resolution: 视频分辨率 ("480p" 或 "720p") video_length: 视频时长(秒) pool_candidates: 候选 Token 池（按优先级） Returns: TokenInfo 对象或 None（无可用 token） """ # 确定首选池 requires_super = resolution == "720p" or video_length > 6 primary_pool = SUPER_POOL_NAME if requires_super else BASIC_POOL_NAME if pool_candidates: ordered_pools = list(pool_candidates) if primary_pool in ordered_pools: ordered_pools.remove(primary_pool) ordered_pools.insert(0, primary_pool) else: fallback_pool = BASIC_POOL_NAME if requires_super else SUPER_POOL_NAME ordered_pools = [primary_pool, fallback_pool] for idx, pool_name in enumerate(ordered_pools): token_info = self.get_token_info(pool_name) if token_info: if idx == 0: logger.info( f"Video token routing: resolution={resolution}, length={video_length}s -> " f"pool={pool_name} (token={token_info.token[:10]}...)" ) else: logger.info( f"Video token routing: fallback from {ordered_pools[0]} -> {pool_name} " f"(token={token_info.token[:10]}...)" ) return token_info if idx == 0 and requires_super and pool_name == primary_pool: next_pool = ordered_pools[1] if len(ordered_pools) > 1 else None if next_pool: logger.warning( f"Video token routing: {primary_pool} pool has no available token for " f"resolution={resolution}, length={video_length}s. " f"Falling back to {next_pool} pool." ) # 两个池都没有可用 token logger.warning( f"Video token routing: no available token in any pool " f"(resolution={resolution}, length={video_length}s)" ) return None def get_pool_name_for_token(self, token_str: str) -> Optional[str]: """Return pool name for the given token string.""" raw_token = token_str.replace("sso=", "") for pool_name, pool in self.pools.items(): if pool.get(raw_token): return pool_name return None async def consume( self, token_str: str, effort: EffortType = EffortType.LOW ) -> bool: """ 消耗配额（本地预估） Args: token_str: Token 字符串 effort: 消耗力度 Returns: 是否成功 """ raw_token = token_str.replace("sso=", "") for pool in self.pools.values(): token = pool.get(raw_token) if token: old_status = token.status consumed = token.consume(effort) logger.debug( f"Token {raw_token[:10]}...: consumed {consumed} quota, use_count={token.use_count}" ) change_kind = "state" if token.status != old_status else "usage" self._track_token_change(token, pool.name, change_kind) self._schedule_save() return True logger.warning(f"Token {raw_token[:10]}...: not found for consumption") return False async def sync_usage( self, token_str: str, fallback_effort: EffortType = EffortType.LOW, consume_on_fail: bool = True, is_usage: bool = True, ) -> bool: """ 同步 Token 用量 优先从 API 获取最新配额，失败则降级到本地预估 Args: token_str: Token 字符串（可带 sso= 前缀） fallback_effort: 降级时的消耗力度 consume_on_fail: 失败时是否降级扣费 is_usage: 是否记录为一次使用（影响 use_count） Returns: 是否成功 """ raw_token = token_str.replace("sso=", "") # 查找 Token 对象 target_token: Optional[TokenInfo] = None target_pool_name: Optional[str] = None for pool in self.pools.values(): target_token = pool.get(raw_token) if target_token: target_pool_name = pool.name break if not target_token: logger.warning(f"Token {raw_token[:10]}...: not found for sync") return False # 尝试 API 同步 try: usage_service = UsageService() result = await usage_service.get(token_str) if result and "remainingTokens" in result: new_quota = result.get("remainingTokens") if new_quota is None: new_quota = result.get("remainingQueries") if new_quota is None: return False old_quota = target_token.quota old_status = target_token.status target_token.update_quota(new_quota) target_token.record_success(is_usage=is_usage) consumed = max(0, old_quota - new_quota) logger.info( f"Token {raw_token[:10]}...: synced quota " f"{old_quota} -> {new_quota} (consumed: {consumed}, use_count: {target_token.use_count})" ) if target_pool_name: change_kind = "state" if target_token.status != old_status else "usage" self._track_token_change( target_token, target_pool_name, change_kind ) self._schedule_save() return True except Exception as e: if isinstance(e, UpstreamException): status = None if e.details and "status" in e.details: status = e.details["status"] else: status = getattr(e, "status_code", None) if status == 401: await self.record_fail(token_str, status, "rate_limits_auth_failed") logger.warning( f"Token {raw_token[:10]}...: API sync failed, fallback to local ({e})" ) # 降级：本地预估扣费 if consume_on_fail: logger.debug(f"Token {raw_token[:10]}...: using local consumption") return await self.consume(token_str, fallback_effort) else: logger.debug( f"Token {raw_token[:10]}...: sync failed, skipping local consumption" ) return False async def record_fail( self, token_str: str, status_code: int = 401, reason: str = "" ) -> bool: """ 记录 Token 失败 Args: token_str: Token 字符串 status_code: HTTP 状态码 reason: 失败原因 Returns: 是否成功 """ raw_token = token_str.replace("sso=", "") for pool in self.pools.values(): token = pool.get(raw_token) if token: if status_code == 401: threshold = get_config("token.fail_threshold", FAIL_THRESHOLD) try: threshold = int(threshold) except (TypeError, ValueError): threshold = FAIL_THRESHOLD if threshold < 1: threshold = 1 token.record_fail(status_code, reason, threshold=threshold) logger.warning( f"Token {raw_token[:10]}...: recorded {status_code} failure " f"({token.fail_count}/{threshold}) - {reason}" ) self._track_token_change(token, pool.name, "state") self._schedule_save() else: logger.info( f"Token {raw_token[:10]}...: non-auth error ({status_code}) - {reason} (not counted)" ) return True logger.warning(f"Token {raw_token[:10]}...: not found for failure record") return False async def mark_rate_limited(self, token_str: str) -> bool: """ 将 Token 标记为配额耗尽（COOLING） 当 Grok API 返回 429 时调用，将 quota 设为 0 并标记 COOLING， 使该 Token 不再被选中，等待下次 Scheduler 刷新恢复。 Args: token_str: Token 字符串 Returns: 是否成功 """ raw_token = token_str.removeprefix("sso=") for pool in self.pools.values(): token = pool.get(raw_token) if token: old_quota = token.quota token.quota = 0 token.status = TokenStatus.COOLING logger.warning( f"Token {raw_token[:10]}...: marked as rate limited " f"(quota {old_quota} -> 0, status -> cooling)" ) self._track_token_change(token, pool.name, "state") self._schedule_save() return True logger.warning(f"Token {raw_token[:10]}...: not found for rate limit marking") return False # ========== 管理功能 ========== async def add(self, token: str, pool_name: str = "ssoBasic") -> bool: """ 添加 Token Args: token: Token 字符串（不含 sso= 前缀） pool_name: 池名称 Returns: 是否成功 """ if pool_name not in self.pools: self.pools[pool_name] = TokenPool(pool_name) logger.info(f"Pool '{pool_name}': created") pool = self.pools[pool_name] token = token[4:] if token.startswith("sso=") else token if pool.get(token): logger.warning(f"Pool '{pool_name}': token already exists") return False token_info = TokenInfo(token=token, quota=_default_quota_for_pool(pool_name)) pool.add(token_info) self._track_token_change(token_info, pool_name, "state") await self._save(force=True) logger.info(f"Pool '{pool_name}': token added") return True async def mark_asset_clear(self, token: str) -> bool: """记录在线资产清理时间""" raw_token = token[4:] if token.startswith("sso=") else token for pool in self.pools.values(): info = pool.get(raw_token) if info: info.last_asset_clear_at = int(datetime.now().timestamp() * 1000) self._track_token_change(info, pool.name, "state") self._schedule_save() return True return False async def add_tag(self, token: str, tag: str) -> bool: """ 给 Token 添加标签 Args: token: Token 字符串 tag: 标签名称 Returns: 是否成功 """ raw_token = token[4:] if token.startswith("sso=") else token for pool in self.pools.values(): info = pool.get(raw_token) if info: if tag not in info.tags: info.tags.append(tag) self._track_token_change(info, pool.name, "state") self._schedule_save() logger.debug(f"Token {raw_token[:10]}...: added tag '{tag}'") return True return False async def remove_tag(self, token: str, tag: str) -> bool: """ 移除 Token 标签 Args: token: Token 字符串 tag: 标签名称 Returns: 是否成功 """ raw_token = token[4:] if token.startswith("sso=") else token for pool in self.pools.values(): info = pool.get(raw_token) if info: if tag in info.tags: info.tags.remove(tag) self._track_token_change(info, pool.name, "state") self._schedule_save() logger.debug(f"Token {raw_token[:10]}...: removed tag '{tag}'") return True return False async def remove(self, token: str) -> bool: """ 删除 Token Args: token: Token 字符串 Returns: 是否成功 """ for pool_name, pool in self.pools.items(): if pool.remove(token): self._track_token_delete(token) await self._save(force=True) logger.info(f"Pool '{pool_name}': token removed") return True logger.warning("Token not found for removal") return False async def reset_all(self): """重置所有 Token 配额""" count = 0 for pool_name, pool in self.pools.items(): default_quota = _default_quota_for_pool(pool_name) for token in pool: token.reset(default_quota) self._track_token_change(token, pool_name, "state") count += 1 await self._save(force=True) logger.info(f"Reset all: {count} tokens updated") async def reset_token(self, token_str: str) -> bool: """ 重置单个 Token Args: token_str: Token 字符串 Returns: 是否成功 """ raw_token = token_str.replace("sso=", "") for pool in self.pools.values(): token = pool.get(raw_token) if token: default_quota = _default_quota_for_pool(pool.name) token.reset(default_quota) self._track_token_change(token, pool.name, "state") await self._save(force=True) logger.info(f"Token {raw_token[:10]}...: reset completed") return True logger.warning(f"Token {raw_token[:10]}...: not found for reset") return False def get_stats(self) -> Dict[str, dict]: """获取统计信息""" stats = {} for name, pool in self.pools.items(): pool_stats = pool.get_stats() stats[name] = pool_stats.model_dump() return stats def get_pool_tokens(self, pool_name: str = "ssoBasic") -> List[TokenInfo]: """ 获取指定池的所有 Token Args: pool_name: 池名称 Returns: Token 列表 """ pool = self.pools.get(pool_name) if not pool: return [] return pool.list() async def refresh_cooling_tokens(self) -> Dict[str, int]: """ 批量刷新 cooling 状态的 Token 配额 Returns: {"checked": int, "refreshed": int, "recovered": int, "expired": int} """ # 收集需要刷新的 token to_refresh: List[tuple[str, TokenInfo]] = [] for pool in self.pools.values(): if pool.name == SUPER_POOL_NAME: interval_hours = get_config( "token.super_refresh_interval_hours", DEFAULT_SUPER_REFRESH_INTERVAL_HOURS, ) else: interval_hours = get_config( "token.refresh_interval_hours", DEFAULT_REFRESH_INTERVAL_HOURS, ) for token in pool: if token.need_refresh(interval_hours): to_refresh.append((pool.name, token)) if not to_refresh: logger.debug("Refresh check: no tokens need refresh") return {"checked": 0, "refreshed": 0, "recovered": 0, "expired": 0} logger.info(f"Refresh check: found {len(to_refresh)} cooling tokens to refresh") # 批量并发刷新 semaphore = asyncio.Semaphore(DEFAULT_REFRESH_CONCURRENCY) usage_service = UsageService() refreshed = 0 recovered = 0 expired = 0 async def _refresh_one(item: tuple[str, TokenInfo]) -> dict: """刷新单个 token""" _, token_info = item async with semaphore: token_str = token_info.token if token_str.startswith("sso="): token_str = token_str[4:] # 重试逻辑：最多 2 次重试 for retry in range(3): # 0, 1, 2 try: result = await usage_service.get(token_str) if result and "remainingTokens" in result: new_quota = result.get("remainingTokens") if new_quota is None: new_quota = result.get("remainingQueries") if new_quota is None: return {"recovered": False, "expired": False} old_quota = token_info.quota old_status = token_info.status token_info.update_quota(new_quota) token_info.mark_synced() logger.info( f"Token {token_info.token[:10]}...: refreshed " f"{old_quota} -> {new_quota}, status: {old_status} -> {token_info.status}" ) return { "recovered": new_quota > 0 and old_quota == 0, "expired": False, } return {"recovered": False, "expired": False} except Exception as e: error_str = str(e) # 检查是否为 401 错误 if "401" in error_str or "Unauthorized" in error_str: if retry < 2: logger.warning( f"Token {token_info.token[:10]}...: 401 error, " f"retry {retry + 1}/2..." ) await asyncio.sleep(0.5) continue else: # 重试 2 次后仍然 401，标记为 expired logger.error( f"Token {token_info.token[:10]}...: 401 after 2 retries, " f"marking as expired" ) token_info.status = TokenStatus.EXPIRED return {"recovered": False, "expired": True} else: logger.warning( f"Token {token_info.token[:10]}...: refresh failed ({e})" ) return {"recovered": False, "expired": False} return {"recovered": False, "expired": False} # 批量处理 for i in range(0, len(to_refresh), DEFAULT_REFRESH_BATCH_SIZE): batch = to_refresh[i : i + DEFAULT_REFRESH_BATCH_SIZE] results = await asyncio.gather(*[_refresh_one(t) for t in batch]) refreshed += len(batch) recovered += sum(r["recovered"] for r in results) expired += sum(r["expired"] for r in results) # 批次间延迟 if i + DEFAULT_REFRESH_BATCH_SIZE < len(to_refresh): await asyncio.sleep(1) for pool_name, token_info in to_refresh: self._track_token_change(token_info, pool_name, "state") await self._save(force=True) logger.info( f"Refresh completed: " f"checked={len(to_refresh)}, refreshed={refreshed}, " f"recovered={recovered}, expired={expired}" ) return { "checked": len(to_refresh), "refreshed": refreshed, "recovered": recovered, "expired": expired, } # 便捷函数 async def get_token_manager() -> TokenManager: """获取 TokenManager 单例""" return await TokenManager.get_instance() __all__ = ["TokenManager", "get_token_manager"]