"""全局状态管理""" import time from collections import deque from dataclasses import dataclass from typing import Optional, List, Dict from pathlib import Path from ..config import TOKEN_PATH from ..credential import quota_manager, CredentialStatus from .account import Account from .persistence import load_accounts, save_accounts from .quota_cache import get_quota_cache from .account_selector import get_account_selector, SelectionStrategy from .quota_scheduler import get_quota_scheduler @dataclass class RequestLog: """请求日志""" id: str timestamp: float method: str path: str model: str account_id: Optional[str] status: int duration_ms: float tokens_in: int = 0 tokens_out: int = 0 error: Optional[str] = None class ProxyState: """全局状态管理""" def __init__(self): self.accounts: List[Account] = [] self.request_logs: deque = deque(maxlen=1000) self.total_requests: int = 0 self.total_errors: int = 0 self.session_locks: Dict[str, str] = {} self.session_timestamps: Dict[str, float] = {} self.start_time: float = time.time() self._load_accounts() def _load_accounts(self): """从配置文件加载账号""" saved = load_accounts() if saved: for acc_data in saved: # 验证 token 文件存在 if Path(acc_data.get("token_path", "")).exists(): self.accounts.append(Account( id=acc_data["id"], name=acc_data["name"], token_path=acc_data["token_path"], enabled=acc_data.get("enabled", True), auto_disabled=acc_data.get("auto_disabled", False), )) print(f"[State] 从配置加载 {len(self.accounts)} 个账号") # 如果没有账号，尝试添加默认账号 if not self.accounts and TOKEN_PATH.exists(): self.accounts.append(Account( id="default", name="默认账号", token_path=str(TOKEN_PATH) )) self._save_accounts() def _save_accounts(self): """保存账号到配置文件""" accounts_data = [ { "id": acc.id, "name": acc.name, "token_path": acc.token_path, "enabled": acc.enabled, "auto_disabled": getattr(acc, "auto_disabled", False), } for acc in self.accounts ] save_accounts(accounts_data) def get_available_account(self, session_id: Optional[str] = None) -> Optional[Account]: """获取可用账号（支持会话粘性和智能选择）""" quota_manager.cleanup_expired() selector = get_account_selector() has_priority = bool(selector.get_priority_accounts()) use_session_sticky = bool(session_id) and not has_priority and selector.strategy != SelectionStrategy.RANDOM # 会话粘性 if use_session_sticky and session_id in self.session_locks: account_id = self.session_locks[session_id] ts = self.session_timestamps.get(session_id, 0) if time.time() - ts < 60: for acc in self.accounts: if acc.id == account_id and acc.is_available(): self.session_timestamps[session_id] = time.time() return acc # 使用 AccountSelector 选择账号 account = selector.select(self.accounts, session_id) if account and use_session_sticky: self.session_locks[session_id] = account.id self.session_timestamps[session_id] = time.time() # 标记为活跃账号，便于额度调度器定期更新 if account: try: get_quota_scheduler().mark_active(account.id) except Exception: pass return account def mark_account_used(self, account_id: str) -> None: """标记账号被使用""" scheduler = get_quota_scheduler() scheduler.mark_active(account_id) for acc in self.accounts: if acc.id == account_id: acc.last_used = time.time() break def get_next_available_account(self, exclude_id: str) -> Optional[Account]: """获取下一个可用账号（排除指定账号）""" available = [a for a in self.accounts if a.is_available() and a.id != exclude_id] if not available: return None account = min(available, key=lambda a: a.request_count) try: get_quota_scheduler().mark_active(account.id) except Exception: pass return account def mark_rate_limited(self, account_id: str, duration_seconds: int = 60): """标记账号限流""" for acc in self.accounts: if acc.id == account_id: acc.mark_quota_exceeded("Rate limited") break def mark_quota_exceeded(self, account_id: str, reason: str = "Quota exceeded"): """标记账号配额超限""" for acc in self.accounts: if acc.id == account_id: acc.mark_quota_exceeded(reason) break async def refresh_account_token(self, account_id: str) -> tuple: """刷新指定账号的 token""" for acc in self.accounts: if acc.id == account_id: return await acc.refresh_token() return False, "账号不存在" async def refresh_expiring_tokens(self) -> List[dict]: """刷新所有即将过期的 token""" results = [] for acc in self.accounts: if acc.enabled and acc.is_token_expiring_soon(10): success, msg = await acc.refresh_token() results.append({ "account_id": acc.id, "success": success, "message": msg }) return results def add_log(self, log: RequestLog): """添加请求日志""" self.request_logs.append(log) self.total_requests += 1 if log.error: self.total_errors += 1 def get_stats(self) -> dict: """获取统计信息""" uptime = time.time() - self.start_time # 获取额度汇总 quota_cache = get_quota_cache() quota_summary = quota_cache.get_summary() # 获取选择器状态 selector = get_account_selector() selector_status = selector.get_status() # 获取调度器状态 scheduler = get_quota_scheduler() scheduler_status = scheduler.get_status() return { "uptime_seconds": int(uptime), "total_requests": self.total_requests, "total_errors": self.total_errors, "error_rate": f"{(self.total_errors / max(1, self.total_requests) * 100):.1f}%", "accounts_total": len(self.accounts), "accounts_available": len([a for a in self.accounts if a.is_available()]), "accounts_cooldown": len([a for a in self.accounts if a.status == CredentialStatus.COOLDOWN]), "recent_logs": len(self.request_logs), # 新增字段 "quota_summary": quota_summary, "selector": selector_status, "scheduler": scheduler_status, } def get_accounts_status(self) -> List[dict]: """获取所有账号状态""" return [acc.get_status_info() for acc in self.accounts] def get_accounts_summary(self) -> dict: """获取账号汇总统计""" quota_cache = get_quota_cache() selector = get_account_selector() scheduler = get_quota_scheduler() total_balance = 0.0 total_usage = 0.0 total_limit = 0.0 available_count = 0 cooldown_count = 0 unhealthy_count = 0 disabled_count = 0 for acc in self.accounts: if not acc.enabled: disabled_count += 1 elif acc.status == CredentialStatus.COOLDOWN: cooldown_count += 1 elif acc.status == CredentialStatus.UNHEALTHY: unhealthy_count += 1 elif acc.is_available(): available_count += 1 quota = quota_cache.get(acc.id) if quota and not quota.has_error(): total_balance += quota.balance total_usage += quota.current_usage total_limit += quota.usage_limit last_refresh = scheduler.get_last_full_refresh() last_refresh_ago = None if last_refresh: seconds_ago = time.time() - last_refresh if seconds_ago < 60: last_refresh_ago = f"{int(seconds_ago)}秒前" elif seconds_ago < 3600: last_refresh_ago = f"{int(seconds_ago / 60)}分钟前" else: last_refresh_ago = f"{int(seconds_ago / 3600)}小时前" return { "total_accounts": len(self.accounts), "available_accounts": available_count, "cooldown_accounts": cooldown_count, "unhealthy_accounts": unhealthy_count, "disabled_accounts": disabled_count, "total_balance": round(total_balance, 2), "total_usage": round(total_usage, 2), "total_limit": round(total_limit, 2), "last_refresh": last_refresh_ago, "last_refresh_timestamp": last_refresh, "strategy": selector.strategy.value, "priority_accounts": selector.get_priority_accounts(), } def get_valid_account_ids(self) -> set: """获取所有有效账号ID集合""" return {acc.id for acc in self.accounts if acc.enabled} # 全局状态实例 state = ProxyState()