"""Token 刷新管理模块 提供 Token 批量刷新的管理功能，包括： - 刷新进度跟踪 - 并发控制 - 重试机制配置 - 全局锁防止重复刷新 - Token 过期检测和自动刷新 - 指数退避重试策略 """ import time import asyncio from dataclasses import dataclass, field, asdict from typing import Optional, Dict, Any, List, Tuple, Callable, TYPE_CHECKING from threading import Lock if TYPE_CHECKING: from .account import Account @dataclass class RefreshProgress: """刷新进度信息 用于跟踪批量 Token 刷新操作的进度状态。 Attributes: total: 需要刷新的账号总数 completed: 已完成处理的账号数（包括成功和失败） success: 刷新成功的账号数 failed: 刷新失败的账号数 current_account: 当前正在处理的账号ID status: 刷新状态 - running(进行中), completed(已完成), error(出错) started_at: 刷新开始时间戳 message: 状态消息，用于显示当前操作或错误信息 """ total: int = 0 completed: int = 0 success: int = 0 failed: int = 0 current_account: Optional[str] = None status: str = "running" # running, completed, error started_at: float = field(default_factory=time.time) message: Optional[str] = None def to_dict(self) -> Dict[str, Any]: """转换为字典格式 Returns: 包含所有进度信息的字典 """ return asdict(self) @property def progress_percent(self) -> float: """计算完成百分比 Returns: 完成百分比（0-100） """ if self.total == 0: return 0.0 return round((self.completed / self.total) * 100, 2) @property def elapsed_seconds(self) -> float: """计算已用时间（秒） Returns: 从开始到现在的秒数 """ return time.time() - self.started_at def is_running(self) -> bool: """检查是否正在运行 Returns: True 表示正在运行 """ return self.status == "running" def is_completed(self) -> bool: """检查是否已完成 Returns: True 表示已完成（成功或出错） """ return self.status in ("completed", "error") @dataclass class RefreshConfig: """刷新配置 控制 Token 刷新行为的配置参数。 Attributes: max_retries: 单个账号刷新失败时的最大重试次数 retry_base_delay: 重试基础延迟时间（秒），实际延迟会指数增长 concurrency: 并发刷新的账号数量 token_refresh_before_expiry: Token 过期前多少秒开始刷新（默认5分钟） auto_refresh_interval: 自动刷新检查间隔（秒） """ max_retries: int = 3 retry_base_delay: float = 1.0 concurrency: int = 3 token_refresh_before_expiry: int = 300 # 5分钟 auto_refresh_interval: int = 60 # 1分钟 def to_dict(self) -> Dict[str, Any]: """转换为字典格式 Returns: 包含所有配置项的字典 """ return asdict(self) @classmethod def from_dict(cls, data: Dict[str, Any]) -> 'RefreshConfig': """从字典创建配置实例 Args: data: 配置字典 Returns: RefreshConfig 实例 """ return cls( max_retries=data.get("max_retries", 3), retry_base_delay=data.get("retry_base_delay", 1.0), concurrency=data.get("concurrency", 3), token_refresh_before_expiry=data.get("token_refresh_before_expiry", 300), auto_refresh_interval=data.get("auto_refresh_interval", 60) ) def validate(self) -> bool: """验证配置有效性 Returns: True 表示配置有效 Raises: ValueError: 配置值无效时抛出 """ if self.max_retries < 0: raise ValueError("max_retries 不能为负数") if self.retry_base_delay <= 0: raise ValueError("retry_base_delay 必须大于0") if self.concurrency < 1: raise ValueError("concurrency 必须至少为1") if self.token_refresh_before_expiry < 0: raise ValueError("token_refresh_before_expiry 不能为负数") if self.auto_refresh_interval < 1: raise ValueError("auto_refresh_interval 必须至少为1秒") return True class RefreshManager: """Token 刷新管理器 管理 Token 批量刷新操作，提供： - 全局锁机制防止重复刷新 - 进度跟踪 - 配置管理 - 自动 Token 刷新定时器 使用示例: manager = get_refresh_manager() if not manager.is_refreshing(): # 开始刷新操作 pass """ def __init__(self, config: Optional[RefreshConfig] = None): """初始化刷新管理器 Args: config: 刷新配置，None 则使用默认配置 """ # 配置 self._config = config or RefreshConfig() # 线程锁（用于同步访问状态） self._lock = Lock() # 异步锁（用于防止并发刷新操作） self._async_lock = asyncio.Lock() # 刷新状态 self._is_refreshing: bool = False self._progress: Optional[RefreshProgress] = None # 上次刷新完成时间 self._last_refresh_time: Optional[float] = None # 自动刷新定时器 self._auto_refresh_task: Optional[asyncio.Task] = None self._auto_refresh_running: bool = False # 获取账号列表的回调函数 self._accounts_getter: Optional[Callable] = None @property def config(self) -> RefreshConfig: """获取当前配置 Returns: 当前的刷新配置 """ with self._lock: return self._config def is_refreshing(self) -> bool: """检查是否正在刷新 Returns: True 表示正在进行刷新操作 """ with self._lock: return self._is_refreshing def get_progress(self) -> Optional[RefreshProgress]: """获取当前刷新进度 Returns: 当前进度信息，如果没有进行中的刷新则返回 None """ with self._lock: return self._progress def get_progress_dict(self) -> Optional[Dict[str, Any]]: """获取当前刷新进度（字典格式） Returns: 进度信息字典，如果没有进行中的刷新则返回 None """ with self._lock: if self._progress is None: return None return self._progress.to_dict() def update_config(self, **kwargs) -> None: """更新配置参数 支持的参数: max_retries: 最大重试次数 retry_base_delay: 重试基础延迟 concurrency: 并发数 token_refresh_before_expiry: Token 过期前刷新时间 auto_refresh_interval: 自动刷新检查间隔 Args: **kwargs: 要更新的配置项 Raises: ValueError: 配置值无效时抛出 """ with self._lock: # 创建新配置 new_config = RefreshConfig( max_retries=kwargs.get("max_retries", self._config.max_retries), retry_base_delay=kwargs.get("retry_base_delay", self._config.retry_base_delay), concurrency=kwargs.get("concurrency", self._config.concurrency), token_refresh_before_expiry=kwargs.get( "token_refresh_before_expiry", self._config.token_refresh_before_expiry ), auto_refresh_interval=kwargs.get( "auto_refresh_interval", self._config.auto_refresh_interval ) ) # 验证配置 new_config.validate() # 应用新配置 self._config = new_config def _start_refresh(self, total: int, message: Optional[str] = None) -> RefreshProgress: """开始刷新操作（内部方法） Args: total: 需要刷新的账号总数 message: 初始状态消息 Returns: 新创建的进度对象 """ with self._lock: self._is_refreshing = True self._progress = RefreshProgress( total=total, completed=0, success=0, failed=0, current_account=None, status="running", started_at=time.time(), message=message or "开始刷新" ) return self._progress def _update_progress( self, current_account: Optional[str] = None, success: bool = False, failed: bool = False, message: Optional[str] = None ) -> None: """更新刷新进度（内部方法） Args: current_account: 当前处理的账号ID success: 是否成功完成一个账号 failed: 是否失败一个账号 message: 状态消息 """ with self._lock: if self._progress is None: return if current_account is not None: self._progress.current_account = current_account if success: self._progress.success += 1 self._progress.completed += 1 elif failed: self._progress.failed += 1 self._progress.completed += 1 if message is not None: self._progress.message = message def _finish_refresh(self, status: str = "completed", message: Optional[str] = None) -> None: """完成刷新操作（内部方法） Args: status: 最终状态 - completed 或 error message: 最终状态消息 """ with self._lock: self._is_refreshing = False self._last_refresh_time = time.time() if self._progress is not None: self._progress.status = status self._progress.current_account = None if message is not None: self._progress.message = message elif status == "completed": self._progress.message = ( f"刷新完成: 成功 {self._progress.success}, " f"失败 {self._progress.failed}" ) def get_last_refresh_time(self) -> Optional[float]: """获取上次刷新完成时间 Returns: 上次刷新完成的时间戳，如果从未刷新则返回 None """ with self._lock: return self._last_refresh_time def get_status(self) -> Dict[str, Any]: """获取管理器状态 Returns: 包含管理器状态信息的字典 """ with self._lock: return { "is_refreshing": self._is_refreshing, "progress": self._progress.to_dict() if self._progress else None, "last_refresh_time": self._last_refresh_time, "config": self._config.to_dict() } async def acquire_refresh_lock(self) -> bool: """尝试获取刷新锁 用于在开始刷新操作前获取异步锁，防止并发刷新。 Returns: True 表示成功获取锁，False 表示已有刷新在进行 """ if self._async_lock.locked(): return False await self._async_lock.acquire() return True def release_refresh_lock(self) -> None: """释放刷新锁 在刷新操作完成后调用，释放异步锁。 """ if self._async_lock.locked(): self._async_lock.release() def should_refresh_token(self, account: 'Account') -> bool: """判断是否需要刷新 Token 检查账号的 Token 是否即将过期（过期前5分钟）或已过期。 Args: account: 账号对象 Returns: True 表示需要刷新 Token """ creds = account.get_credentials() if creds is None: return True # 无法获取凭证，需要刷新 # 检查是否已过期或即将过期 minutes_before = self._config.token_refresh_before_expiry // 60 return creds.is_expired() or creds.is_expiring_soon(minutes=minutes_before) async def refresh_token_if_needed(self, account: 'Account') -> Tuple[bool, str]: """如果需要则刷新 Token 检查账号 Token 状态，如果即将过期或已过期则刷新。 Args: account: 账号对象 Returns: (success, message) 元组 - success: True 表示 Token 有效（无需刷新或刷新成功） - message: 状态消息 """ if not self.should_refresh_token(account): return True, "Token 有效，无需刷新" print(f"[RefreshManager] 账号 {account.id} Token 即将过期，开始刷新...") success, result = await account.refresh_token() if success: print(f"[RefreshManager] 账号 {account.id} Token 刷新成功") return True, "Token 刷新成功" else: print(f"[RefreshManager] 账号 {account.id} Token 刷新失败: {result}") return False, f"Token 刷新失败: {result}" async def refresh_account_with_token( self, account: 'Account', get_quota_func: Optional[Callable] = None ) -> Tuple[bool, str]: """刷新单个账号（先刷新 Token，再获取额度） Args: account: 账号对象 get_quota_func: 获取额度的异步函数，接受 account 参数 Returns: (success, message) 元组 """ # 1. 先刷新 Token（如果需要） token_success, token_msg = await self.refresh_token_if_needed(account) if not token_success: return False, token_msg # 2. 获取额度（如果提供了获取函数） if get_quota_func: try: quota_success, quota_result = await get_quota_func(account) if quota_success: return True, "刷新成功" else: error_msg = quota_result.get("error", "Unknown error") if isinstance(quota_result, dict) else str(quota_result) return False, f"获取额度失败: {error_msg}" except Exception as e: return False, f"获取额度异常: {str(e)}" return True, token_msg async def retry_with_backoff( self, func: Callable, *args, max_retries: Optional[int] = None, **kwargs ) -> Tuple[bool, Any]: """带指数退避的重试 执行异步函数，失败时使用指数退避策略重试。 Args: func: 要执行的异步函数 *args: 传递给函数的位置参数 max_retries: 最大重试次数，None 则使用配置值 **kwargs: 传递给函数的关键字参数 Returns: (success, result) 元组 - success: True 表示执行成功 - result: 成功时为函数返回值，失败时为错误信息 """ retries = max_retries if max_retries is not None else self._config.max_retries base_delay = self._config.retry_base_delay last_error = None for attempt in range(retries + 1): try: result = await func(*args, **kwargs) # 检查返回值格式 if isinstance(result, tuple) and len(result) == 2: success, data = result if success: return True, data else: last_error = data # 检查是否是 429 错误 if self._is_rate_limit_error(data): delay = self._get_rate_limit_delay(attempt, base_delay) else: delay = base_delay * (2 ** attempt) else: # 函数返回非元组，视为成功 return True, result except Exception as e: last_error = str(e) delay = base_delay * (2 ** attempt) # 如果还有重试机会，等待后重试 if attempt < retries: print(f"[RefreshManager] 第 {attempt + 1} 次尝试失败，{delay:.1f}秒后重试...") await asyncio.sleep(delay) return False, last_error def _is_rate_limit_error(self, error: Any) -> bool: """检查是否是限流错误（429） Args: error: 错误信息 Returns: True 表示是限流错误 """ if isinstance(error, str): return "429" in error or "rate limit" in error.lower() or "请求过于频繁" in error return False def _get_rate_limit_delay(self, attempt: int, base_delay: float) -> float: """获取限流错误的等待时间 429 错误使用更长的等待时间。 Args: attempt: 当前尝试次数（从0开始） base_delay: 基础延迟 Returns: 等待时间（秒） """ # 429 错误使用 3 倍的基础延迟 return base_delay * 3 * (2 ** attempt) async def refresh_all_with_token( self, accounts: List['Account'], get_quota_func: Optional[Callable] = None, skip_disabled: bool = True, skip_error: bool = True ) -> RefreshProgress: """刷新所有账号（先刷新 Token，再获取额度） 使用全局锁防止并发刷新，支持进度跟踪。 Args: accounts: 账号列表 get_quota_func: 获取额度的异步函数 skip_disabled: 是否跳过已禁用的账号 skip_error: 是否跳过已处于错误状态的账号 Returns: 刷新进度信息 """ # 尝试获取锁 if not await self.acquire_refresh_lock(): # 已有刷新在进行 progress = self.get_progress() if progress: return progress # 返回一个错误状态的进度 return RefreshProgress( total=0, status="error", message="刷新操作正在进行中" ) try: # 过滤账号 accounts_to_refresh = [] for acc in accounts: if skip_disabled and not acc.enabled: continue if skip_error and acc.status.value in ("unhealthy", "suspended"): continue accounts_to_refresh.append(acc) total = len(accounts_to_refresh) # 开始刷新 self._start_refresh(total, f"开始刷新 {total} 个账号") if total == 0: self._finish_refresh("completed", "没有需要刷新的账号") return self.get_progress() # 使用信号量控制并发 semaphore = asyncio.Semaphore(self._config.concurrency) async def refresh_one(account: 'Account'): async with semaphore: self._update_progress( current_account=account.id, message=f"正在刷新: {account.name}" ) # 使用重试机制刷新 success, result = await self.retry_with_backoff( self.refresh_account_with_token, account, get_quota_func ) if success: self._update_progress(success=True) else: self._update_progress(failed=True) return success, result # 并发执行 tasks = [refresh_one(acc) for acc in accounts_to_refresh] await asyncio.gather(*tasks, return_exceptions=True) # 完成 self._finish_refresh("completed") return self.get_progress() except Exception as e: self._finish_refresh("error", f"刷新异常: {str(e)}") return self.get_progress() finally: self.release_refresh_lock() def _is_auth_error(self, error: Any) -> bool: """检查是否是认证错误（401） Args: error: 错误信息 Returns: True 表示是认证错误 """ if isinstance(error, str): return "401" in error or "unauthorized" in error.lower() or "凭证已过期" in error or "需要重新登录" in error return False async def execute_with_auth_retry( self, account: 'Account', func: Callable, *args, **kwargs ) -> Tuple[bool, Any]: """执行操作，遇到 401 错误时自动刷新 Token 并重试 Args: account: 账号对象 func: 要执行的异步函数 *args: 传递给函数的位置参数 **kwargs: 传递给函数的关键字参数 Returns: (success, result) 元组 """ try: result = await func(*args, **kwargs) # 检查返回值 if isinstance(result, tuple) and len(result) == 2: success, data = result if success: return True, data # 检查是否是 401 错误 if self._is_auth_error(data): print(f"[RefreshManager] 账号 {account.id} 遇到 401 错误，尝试刷新 Token...") # 刷新 Token refresh_success, refresh_msg = await account.refresh_token() if refresh_success: print(f"[RefreshManager] Token 刷新成功，重试请求...") # 重试原请求 retry_result = await func(*args, **kwargs) if isinstance(retry_result, tuple) and len(retry_result) == 2: return retry_result return True, retry_result else: return False, f"Token 刷新失败: {refresh_msg}" return False, data return True, result except Exception as e: error_str = str(e) # 检查异常是否是 401 错误 if self._is_auth_error(error_str): print(f"[RefreshManager] 账号 {account.id} 遇到 401 异常，尝试刷新 Token...") refresh_success, refresh_msg = await account.refresh_token() if refresh_success: print(f"[RefreshManager] Token 刷新成功，重试请求...") try: retry_result = await func(*args, **kwargs) if isinstance(retry_result, tuple) and len(retry_result) == 2: return retry_result return True, retry_result except Exception as retry_e: return False, f"重试失败: {str(retry_e)}" else: return False, f"Token 刷新失败: {refresh_msg}" return False, error_str def set_accounts_getter(self, getter: Callable) -> None: """设置获取账号列表的回调函数 Args: getter: 返回账号列表的可调用对象 """ self._accounts_getter = getter def _get_accounts(self) -> List['Account']: """获取账号列表""" if self._accounts_getter: return self._accounts_getter() return [] async def start_auto_refresh(self) -> None: """启动自动 Token 刷新定时器 定期检查所有账号的 Token 状态，自动刷新即将过期的 Token。 启动前会清除已存在的定时器，防止重复启动。 """ # 先停止已存在的定时器 await self.stop_auto_refresh() self._auto_refresh_running = True self._auto_refresh_task = asyncio.create_task(self._auto_refresh_loop()) print(f"[RefreshManager] 自动 Token 刷新定时器已启动，检查间隔: {self._config.auto_refresh_interval}秒") async def stop_auto_refresh(self) -> None: """停止自动 Token 刷新定时器""" self._auto_refresh_running = False if self._auto_refresh_task: self._auto_refresh_task.cancel() try: await self._auto_refresh_task except asyncio.CancelledError: pass self._auto_refresh_task = None print("[RefreshManager] 自动 Token 刷新定时器已停止") def is_auto_refresh_running(self) -> bool: """检查自动刷新定时器是否在运行 Returns: True 表示定时器正在运行 """ return self._auto_refresh_running and self._auto_refresh_task is not None async def _auto_refresh_loop(self) -> None: """自动刷新循环 定期检查所有账号的 Token 状态，刷新即将过期的 Token。 跳过已禁用或错误状态的账号，单个失败不影响其他账号。 """ while self._auto_refresh_running: try: await asyncio.sleep(self._config.auto_refresh_interval) if not self._auto_refresh_running: break accounts = self._get_accounts() if not accounts: continue # 检查需要刷新的账号 accounts_to_refresh = [] for account in accounts: # 跳过已禁用的账号 if not account.enabled: continue # 跳过错误状态的账号 if hasattr(account, 'status') and account.status.value in ("unhealthy", "suspended", "disabled"): continue # 检查是否需要刷新 Token if self.should_refresh_token(account): accounts_to_refresh.append(account) if accounts_to_refresh: print(f"[RefreshManager] 发现 {len(accounts_to_refresh)} 个账号需要刷新 Token") # 逐个刷新，单个失败不影响其他 for account in accounts_to_refresh: try: success, message = await self.refresh_token_if_needed(account) if not success: print(f"[RefreshManager] 账号 {account.id} 自动刷新失败: {message}") except Exception as e: print(f"[RefreshManager] 账号 {account.id} 自动刷新异常: {e}") # 继续处理其他账号 except asyncio.CancelledError: break except Exception as e: print(f"[RefreshManager] 自动刷新循环异常: {e}") # 继续运行，不因异常停止 def get_auto_refresh_status(self) -> Dict[str, Any]: """获取自动刷新状态 Returns: 包含自动刷新状态信息的字典 """ return { "running": self.is_auto_refresh_running(), "interval": self._config.auto_refresh_interval, "token_refresh_before_expiry": self._config.token_refresh_before_expiry } # 全局刷新管理器实例 _refresh_manager: Optional[RefreshManager] = None _manager_lock = Lock() def get_refresh_manager() -> RefreshManager: """获取全局刷新管理器实例 使用单例模式，确保全局只有一个刷新管理器实例。 Returns: 全局 RefreshManager 实例 """ global _refresh_manager if _refresh_manager is None: with _manager_lock: # 双重检查锁定 if _refresh_manager is None: _refresh_manager = RefreshManager() return _refresh_manager def reset_refresh_manager() -> None: """重置全局刷新管理器 主要用于测试场景，重置全局实例。 """ global _refresh_manager with _manager_lock: _refresh_manager = None