|
|
|
|
|
|
| """
|
| Token 池管理器 - 基于数据库的 Token 轮询和健康检查
|
|
|
| 核心功能：
|
| 1. Token 轮询机制 - 负载均衡和容错
|
| 2. Z.AI 官方认证接口验证 - 基于 role 字段区分用户类型
|
| 3. Token 健康度监控 - 自动禁用失败 Token
|
| 4. 数据库集成 - 与 TokenDAO 协同工作
|
| """
|
|
|
| import asyncio
|
| import time
|
| from dataclasses import dataclass
|
| from asyncio import Lock
|
| from typing import Dict, List, Optional, Set, Tuple
|
|
|
| import httpx
|
|
|
| from app.core.resin_compat import apply_resin_account_header
|
| from app.core.upstream_urls import build_upstream_url
|
| from app.core.httpx_errors import normalize_httpx_exception
|
| from app.utils.logger import logger, log_exception
|
| from app.utils.utlis import mask_token
|
|
|
|
|
|
|
|
|
|
|
| TOKEN_LOAD_BALANCE_ROUND_ROBIN = "round_robin"
|
| TOKEN_LOAD_BALANCE_LEAST_LOADED = "least_loaded"
|
| TOKEN_LOAD_BALANCE_SMOOTH_WEIGHTED = "smooth_weighted_round_robin"
|
| SUPPORTED_TOKEN_LOAD_BALANCE_STRATEGIES = {
|
| TOKEN_LOAD_BALANCE_ROUND_ROBIN,
|
| TOKEN_LOAD_BALANCE_LEAST_LOADED,
|
| TOKEN_LOAD_BALANCE_SMOOTH_WEIGHTED,
|
| }
|
|
|
|
|
| def normalize_token_load_balance_strategy(strategy: Optional[str]) -> str:
|
| normalized = str(strategy or "").strip().lower()
|
| if normalized in SUPPORTED_TOKEN_LOAD_BALANCE_STRATEGIES:
|
| return normalized
|
| return TOKEN_LOAD_BALANCE_SMOOTH_WEIGHTED
|
|
|
|
|
| @dataclass
|
| class TokenStatus:
|
| """Token 运行时状态（内存中）"""
|
| token: str
|
| token_id: int
|
| token_type: str = "unknown"
|
| priority: int = 0
|
| is_available: bool = True
|
| failure_count: int = 0
|
| cooldown_until: float = 0.0
|
| last_failure_time: float = 0.0
|
| last_success_time: float = 0.0
|
| last_selected_time: float = 0.0
|
| active_requests: int = 0
|
| current_weight: int = 0
|
| total_requests: int = 0
|
| successful_requests: int = 0
|
| db_synced_successful_requests: int = 0
|
| db_synced_failed_requests: int = 0
|
|
|
| @property
|
| def success_rate(self) -> float:
|
| """成功率"""
|
| if self.total_requests == 0:
|
| return 1.0
|
| return self.successful_requests / self.total_requests
|
|
|
| @property
|
| def failed_requests(self) -> int:
|
| """失败次数。"""
|
| return max(0, self.total_requests - self.successful_requests)
|
|
|
| @property
|
| def is_healthy(self) -> bool:
|
| """
|
| Token 健康状态判断
|
|
|
| 健康标准：
|
| 1. 必须是认证用户 Token (token_type = "user")
|
| 2. 当前可用 (is_available = True)
|
| 3. 成功率 >= 50% 或总请求数 <= 3（新 Token 容错）
|
|
|
| 注意：
|
| - guest Token 永远不健康
|
| - unknown Token 永远不健康
|
| """
|
|
|
| if self.token_type != "user":
|
| return False
|
|
|
|
|
| if not self.is_available:
|
| return False
|
|
|
|
|
| if self.total_requests <= 3:
|
| return self.failure_count == 0
|
|
|
|
|
| return self.success_rate >= 0.5
|
|
|
|
|
|
|
|
|
|
|
| class ZAITokenValidator:
|
| """Z.AI Token 验证器（使用官方认证接口）"""
|
|
|
| @staticmethod
|
| def _response_size_bytes(response: httpx.Response) -> int:
|
| content_length = response.headers.get("content-length")
|
| if content_length and content_length.isdigit():
|
| return int(content_length)
|
| return len(response.content or b"")
|
|
|
| @staticmethod
|
| def auth_url() -> str:
|
| return build_upstream_url("/api/v1/auths/")
|
|
|
| @staticmethod
|
| def get_headers(token: str) -> Dict[str, str]:
|
| """构建认证请求头"""
|
| headers = {
|
| "Accept": "*/*",
|
| "Accept-Language": "zh-CN,zh;q=0.9",
|
| "Authorization": f"Bearer {token}",
|
| "Connection": "keep-alive",
|
| "Content-Type": "application/json",
|
| "DNT": "1",
|
| "Referer": "https://chat.z.ai/",
|
| "Sec-Fetch-Dest": "empty",
|
| "Sec-Fetch-Mode": "cors",
|
| "Sec-Fetch-Site": "same-origin",
|
| "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36",
|
| "sec-ch-ua": '"Chromium";v="140", "Not=A?Brand";v="24", "Google Chrome";v="140"',
|
| "sec-ch-ua-mobile": "?0",
|
| "sec-ch-ua-platform": '"Windows"'
|
| }
|
| apply_resin_account_header(headers, token)
|
| return headers
|
|
|
| @classmethod
|
| async def validate_token(cls, token: str) -> Tuple[str, bool, Optional[str]]:
|
| """
|
| 验证 Token 有效性并返回类型
|
|
|
| Args:
|
| token: 待验证的 Token
|
|
|
| Returns:
|
| (token_type, is_valid, error_message)
|
| - token_type: "user" | "guest" | "unknown"
|
| - is_valid: True 表示是有效的认证用户 Token
|
| - error_message: 失败原因（仅在 is_valid=False 时有值）
|
| """
|
| started_at = time.perf_counter()
|
| masked_token = mask_token(token)
|
| try:
|
| async with httpx.AsyncClient(timeout=15.0, trust_env=False) as client:
|
| auth_url = cls.auth_url()
|
| response = await client.get(
|
| auth_url,
|
| headers=cls.get_headers(token)
|
| )
|
|
|
| elapsed_ms = (time.perf_counter() - started_at) * 1000
|
| logger.debug(
|
| "[token.validate] token=%s status=%s bytes=%s elapsed_ms=%.1f",
|
| masked_token,
|
| response.status_code,
|
| cls._response_size_bytes(response),
|
| elapsed_ms,
|
| )
|
|
|
|
|
| return cls._parse_auth_response(response)
|
|
|
| except httpx.TimeoutException as exc:
|
| return (
|
| "unknown",
|
| False,
|
| normalize_httpx_exception(
|
| exc,
|
| method="GET",
|
| url=cls.auth_url(),
|
| context="token.validate",
|
| ),
|
| )
|
| except httpx.ConnectError as exc:
|
| return (
|
| "unknown",
|
| False,
|
| normalize_httpx_exception(
|
| exc,
|
| method="GET",
|
| url=cls.auth_url(),
|
| context="token.validate",
|
| ),
|
| )
|
| except Exception as e:
|
| return (
|
| "unknown",
|
| False,
|
| normalize_httpx_exception(
|
| e,
|
| method="GET",
|
| url=cls.auth_url(),
|
| context="token.validate",
|
| ),
|
| )
|
|
|
| @staticmethod
|
| def _parse_auth_response(response: httpx.Response) -> Tuple[str, bool, Optional[str]]:
|
| """
|
| 解析 Z.AI 认证接口响应
|
|
|
| 响应格式示例：
|
| {
|
| "id": "...",
|
| "email": "user@example.com",
|
| "role": "user" # 或 "guest"
|
| }
|
|
|
| 验证规则：
|
| - role: "user" → 认证用户 Token（有效，可添加）
|
| - role: "guest" → 匿名用户 Token（无效，拒绝添加）
|
| - 其他情况 → 无效 Token
|
| """
|
|
|
| if response.status_code != 200:
|
| return ("unknown", False, f"HTTP {response.status_code}")
|
|
|
| try:
|
| data = response.json()
|
|
|
|
|
| if not isinstance(data, dict):
|
| return ("unknown", False, "无效的响应格式")
|
|
|
|
|
| if "error" in data or "message" in data:
|
| error_msg = data.get("error") or data.get("message", "未知错误")
|
| return ("unknown", False, str(error_msg))
|
|
|
|
|
| role = data.get("role")
|
|
|
| if role == "user":
|
| return ("user", True, None)
|
| elif role == "guest":
|
| return ("guest", False, "匿名用户 Token 不允许添加")
|
| else:
|
| return ("unknown", False, f"未知 role: {role}")
|
|
|
| except (ValueError, Exception) as e:
|
| return ("unknown", False, f"解析响应失败: {str(e)}")
|
|
|
|
|
|
|
|
|
|
|
| class TokenPool:
|
| """Token 池管理器（数据库驱动）"""
|
|
|
| def __init__(
|
| self,
|
| tokens: List[Tuple[int, str, str]],
|
| failure_threshold: int = 3,
|
| recovery_timeout: int = 1800,
|
| strategy: str = TOKEN_LOAD_BALANCE_SMOOTH_WEIGHTED,
|
| ):
|
| """
|
| 初始化 Token 池
|
|
|
| Args:
|
| tokens: Token 列表 [(token_id, token_value, token_type), ...]
|
| failure_threshold: 失败阈值，超过此次数将标记为不可用
|
| recovery_timeout: 恢复超时时间（秒），失败 Token 在此时间后重新尝试
|
| """
|
| self.failure_threshold = failure_threshold
|
| self.recovery_timeout = recovery_timeout
|
| self._lock = Lock()
|
| self._current_index = 0
|
| self.strategy = self._normalize_strategy(strategy)
|
|
|
|
|
| self.token_statuses: Dict[str, TokenStatus] = {}
|
| self.token_id_map: Dict[str, int] = {}
|
| self._user_tokens: List[str] = []
|
|
|
| for token_record in tokens:
|
| if len(token_record) >= 4:
|
| token_id, token_value, token_type, priority = token_record[:4]
|
| else:
|
| token_id, token_value, token_type = token_record[:3]
|
| priority = 0
|
| if token_value and token_value not in self.token_statuses:
|
| self.token_statuses[token_value] = TokenStatus(
|
| token=token_value,
|
| token_id=token_id,
|
| token_type=token_type,
|
| priority=int(priority or 0),
|
| )
|
| self.token_id_map[token_value] = token_id
|
| if token_type == "user":
|
| self._user_tokens.append(token_value)
|
|
|
| if not self.token_statuses:
|
| logger.warning("⚠️ Token 池为空，将依赖匿名模式")
|
|
|
| @staticmethod
|
| def _normalize_strategy(strategy: Optional[str]) -> str:
|
| return normalize_token_load_balance_strategy(strategy)
|
|
|
| @staticmethod
|
| def _priority_weight(priority: int) -> int:
|
| return max(1, int(priority or 0) + 1)
|
|
|
| async def set_strategy(self, strategy: str) -> None:
|
| normalized = self._normalize_strategy(strategy)
|
| async with self._lock:
|
| self.strategy = normalized
|
| for status in self.token_statuses.values():
|
| status.current_weight = 0
|
|
|
| async def release_token_allocation(self, token: str) -> None:
|
| """释放一次已占用但未完成记账的 token 并发槽位。"""
|
| async with self._lock:
|
| status = self.token_statuses.get(token)
|
| if status is None:
|
| return
|
| status.active_requests = max(0, status.active_requests - 1)
|
|
|
| def _select_round_robin(
|
| self,
|
| candidates: List[Tuple[int, str, TokenStatus]],
|
| total_tokens: int,
|
| current_time: float,
|
| ) -> Optional[str]:
|
| if not candidates:
|
| return None
|
| index, token, status = candidates[0]
|
| status.active_requests += 1
|
| status.last_selected_time = current_time
|
| self._current_index = (index + 1) % total_tokens
|
| return token
|
|
|
| def _select_least_loaded(
|
| self,
|
| candidates: List[Tuple[int, str, TokenStatus]],
|
| total_tokens: int,
|
| current_time: float,
|
| ) -> Optional[str]:
|
| if not candidates:
|
| return None
|
| index, token, status = min(
|
| candidates,
|
| key=lambda item: (
|
| item[2].active_requests,
|
| item[2].total_requests,
|
| item[2].last_selected_time,
|
| item[0],
|
| ),
|
| )
|
| status.active_requests += 1
|
| status.last_selected_time = current_time
|
| self._current_index = (index + 1) % total_tokens
|
| return token
|
|
|
| def _select_smooth_weighted(
|
| self,
|
| candidates: List[Tuple[int, str, TokenStatus]],
|
| total_tokens: int,
|
| current_time: float,
|
| ) -> Optional[str]:
|
| if not candidates:
|
| return None
|
|
|
| min_active_requests = min(status.active_requests for _, _, status in candidates)
|
| shortlist = [
|
| (index, token, status)
|
| for index, token, status in candidates
|
| if status.active_requests == min_active_requests
|
| ]
|
|
|
| total_weight = 0
|
| selected: Optional[Tuple[int, str, TokenStatus]] = None
|
|
|
| for candidate in shortlist:
|
| index, token, status = candidate
|
| weight = self._priority_weight(status.priority)
|
| status.current_weight += weight
|
| total_weight += weight
|
|
|
| if selected is None:
|
| selected = candidate
|
| continue
|
|
|
| selected_index, selected_token, selected_status = selected
|
| if (
|
| status.current_weight > selected_status.current_weight
|
| or (
|
| status.current_weight == selected_status.current_weight
|
| and (
|
| status.total_requests < selected_status.total_requests
|
| or (
|
| status.total_requests == selected_status.total_requests
|
| and (
|
| status.last_selected_time < selected_status.last_selected_time
|
| or (
|
| status.last_selected_time == selected_status.last_selected_time
|
| and index < selected_index
|
| )
|
| )
|
| )
|
| )
|
| )
|
| ):
|
| selected = candidate
|
|
|
| if selected is None:
|
| return None
|
|
|
| index, token, status = selected
|
| status.current_weight -= total_weight
|
| status.active_requests += 1
|
| status.last_selected_time = current_time
|
| self._current_index = (index + 1) % total_tokens
|
| return token
|
|
|
| async def get_next_token(self, exclude_tokens: Optional[Set[str]] = None) -> Optional[str]:
|
| """
|
| 获取下一个可用的认证用户 Token（可配置调度算法）
|
|
|
| Returns:
|
| 可用的 Token 字符串，如果没有可用 Token 则返回 None
|
| """
|
| async with self._lock:
|
| if not self._user_tokens:
|
| return None
|
|
|
| excluded = exclude_tokens or set()
|
| current_time = time.time()
|
| total_tokens = len(self._user_tokens)
|
| start_index = self._current_index % total_tokens if total_tokens else 0
|
| candidate_tokens: List[Tuple[int, str, TokenStatus]] = []
|
|
|
| for offset in range(total_tokens):
|
| index = (start_index + offset) % total_tokens
|
| token = self._user_tokens[index]
|
| if token in excluded:
|
| continue
|
|
|
| status = self.token_statuses.get(token)
|
| if status is None:
|
| continue
|
| if self._is_token_available(status, current_time):
|
| candidate_tokens.append((index, token, status))
|
|
|
| if candidate_tokens:
|
| if self.strategy == TOKEN_LOAD_BALANCE_ROUND_ROBIN:
|
| return self._select_round_robin(
|
| candidate_tokens,
|
| total_tokens,
|
| current_time,
|
| )
|
| if self.strategy == TOKEN_LOAD_BALANCE_LEAST_LOADED:
|
| return self._select_least_loaded(
|
| candidate_tokens,
|
| total_tokens,
|
| current_time,
|
| )
|
| return self._select_smooth_weighted(
|
| candidate_tokens,
|
| total_tokens,
|
| current_time,
|
| )
|
|
|
| logger.warning("⚠️ 没有可用的认证用户 Token")
|
| return None
|
|
|
| def _is_token_available(self, status: TokenStatus, current_time: float) -> bool:
|
| """检查 token 当前是否可用于轮询，必要时自动从熔断状态恢复。"""
|
| if status.token_type != "user":
|
| return False
|
|
|
| if status.is_available:
|
| return True
|
|
|
| if status.cooldown_until and current_time >= status.cooldown_until:
|
| status.is_available = True
|
| status.failure_count = 0
|
| status.cooldown_until = 0.0
|
| logger.info("✅ Token 熔断冷却结束，恢复可用: %s...", status.token[:20])
|
| return True
|
|
|
| return False
|
|
|
| async def mark_token_success(self, token: str):
|
| """标记 Token 使用成功"""
|
| async with self._lock:
|
| if token in self.token_statuses:
|
| status = self.token_statuses[token]
|
| status.total_requests += 1
|
| status.successful_requests += 1
|
| status.active_requests = max(0, status.active_requests - 1)
|
| status.last_success_time = time.time()
|
| status.failure_count = 0
|
| status.cooldown_until = 0.0
|
| status.is_available = True
|
|
|
| async def mark_token_failure(self, token: str, error: Exception = None):
|
| """标记 Token 使用失败"""
|
| async with self._lock:
|
| if token in self.token_statuses:
|
| status = self.token_statuses[token]
|
| current_time = time.time()
|
| status.total_requests += 1
|
| status.active_requests = max(0, status.active_requests - 1)
|
| status.failure_count += 1
|
| status.last_failure_time = current_time
|
|
|
| if status.failure_count >= self.failure_threshold:
|
| status.is_available = False
|
| status.cooldown_until = current_time + self.recovery_timeout
|
| logger.warning(
|
| "Token 已熔断: %s... (连续失败 %s 次, 冷却 %ss)",
|
| token[:20],
|
| status.failure_count,
|
| self.recovery_timeout,
|
| )
|
|
|
| async def record_token_success(self, token: str, dao=None):
|
| """标记成功，并默认异步批量同步数据库统计。"""
|
| await self.mark_token_success(token)
|
|
|
| token_id = self.get_token_id(token)
|
| if token_id is None:
|
| return
|
|
|
| if dao is None:
|
| from app.services.token_dao import get_token_dao
|
| from app.services.token_stats_writer import get_token_stats_writer
|
|
|
| writer = get_token_stats_writer()
|
| enqueued = writer.enqueue(
|
| token=token,
|
| token_id=token_id,
|
| successful_requests=1,
|
| )
|
| if enqueued:
|
| return
|
| dao = get_token_dao()
|
|
|
| try:
|
| await dao.record_success(token_id)
|
| except Exception as e:
|
| log_exception(logger, "❌ 同步 Token 成功统计失败")
|
| return
|
|
|
| await self.mark_stats_synced(token=token, successful_requests=1)
|
|
|
| async def record_token_failure(self, token: str, error: Exception = None, dao=None):
|
| """标记失败，并默认异步批量同步数据库统计。"""
|
| await self.mark_token_failure(token, error)
|
|
|
| token_id = self.get_token_id(token)
|
| if token_id is None:
|
| return
|
|
|
| if dao is None:
|
| from app.services.token_dao import get_token_dao
|
| from app.services.token_stats_writer import get_token_stats_writer
|
|
|
| writer = get_token_stats_writer()
|
| enqueued = writer.enqueue(
|
| token=token,
|
| token_id=token_id,
|
| failed_requests=1,
|
| )
|
| if enqueued:
|
| return
|
| dao = get_token_dao()
|
|
|
| try:
|
| await dao.record_failure(token_id)
|
| except Exception as e:
|
| log_exception(logger, "❌ 同步 Token 失败统计失败")
|
| return
|
|
|
| await self.mark_stats_synced(token=token, failed_requests=1)
|
|
|
| async def mark_stats_synced(
|
| self,
|
| *,
|
| token: str,
|
| successful_requests: int = 0,
|
| failed_requests: int = 0,
|
| ) -> None:
|
| async with self._lock:
|
| if token not in self.token_statuses:
|
| return
|
| status = self.token_statuses[token]
|
| status.db_synced_successful_requests += max(0, successful_requests)
|
| status.db_synced_failed_requests += max(0, failed_requests)
|
|
|
| def get_token_id(self, token: str) -> Optional[int]:
|
| """获取 Token 的数据库 ID"""
|
| return self.token_id_map.get(token)
|
|
|
| async def get_pool_status(self) -> Dict:
|
| """获取 Token 池状态信息"""
|
| async with self._lock:
|
| current_time = time.time()
|
| available_count = sum(
|
| 1
|
| for token in self._user_tokens
|
| if (status := self.token_statuses.get(token))
|
| and self._is_token_available(status, current_time)
|
| )
|
| total_count = len(self.token_statuses)
|
| healthy_count = sum(1 for status in self.token_statuses.values() if status.is_healthy)
|
|
|
|
|
| user_count = sum(1 for s in self.token_statuses.values() if s.token_type == "user")
|
| guest_count = sum(1 for s in self.token_statuses.values() if s.token_type == "guest")
|
| unknown_count = sum(1 for s in self.token_statuses.values() if s.token_type == "unknown")
|
|
|
| status_info = {
|
| "strategy": self.strategy,
|
| "total_tokens": total_count,
|
| "available_tokens": available_count,
|
| "unavailable_tokens": total_count - available_count,
|
| "healthy_tokens": healthy_count,
|
| "unhealthy_tokens": total_count - healthy_count,
|
| "user_tokens": user_count,
|
| "guest_tokens": guest_count,
|
| "unknown_tokens": unknown_count,
|
| "current_index": self._current_index,
|
| "tokens": []
|
| }
|
|
|
| for token, status in self.token_statuses.items():
|
| status_info["tokens"].append({
|
| "token": f"{token[:10]}...{token[-10:]}",
|
| "token_id": status.token_id,
|
| "token_type": status.token_type,
|
| "is_available": status.is_available,
|
| "failure_count": status.failure_count,
|
| "cooldown_until": status.cooldown_until,
|
| "priority": status.priority,
|
| "selection_weight": self._priority_weight(status.priority),
|
| "active_requests": status.active_requests,
|
| "current_weight": status.current_weight,
|
| "success_count": status.successful_requests,
|
| "success_rate": f"{status.success_rate:.2%}",
|
| "total_requests": status.total_requests,
|
| "is_healthy": status.is_healthy,
|
| "last_failure_time": status.last_failure_time,
|
| "last_success_time": status.last_success_time,
|
| "last_selected_time": status.last_selected_time,
|
| })
|
|
|
| return status_info
|
|
|
| async def update_token_type(self, token: str, token_type: str):
|
| """更新 Token 类型（用于健康检查后更新）"""
|
| async with self._lock:
|
| if token in self.token_statuses:
|
| old_type = self.token_statuses[token].token_type
|
| self.token_statuses[token].token_type = token_type
|
|
|
| if old_type != token_type:
|
| logger.debug("更新 Token 类型: %s... %s -> %s", token[:20], old_type, token_type)
|
|
|
| async def health_check_token(self, token: str) -> bool:
|
| """
|
| 异步健康检查单个 Token（使用 Z.AI 官方认证接口）
|
|
|
| Args:
|
| token: 要检查的 Token
|
|
|
| Returns:
|
| Token 是否健康（True = 有效的认证用户 Token）
|
| """
|
| token_type, is_valid, error_message = await ZAITokenValidator.validate_token(token)
|
|
|
|
|
| await self.update_token_type(token, token_type)
|
|
|
|
|
| if is_valid:
|
| await self.record_token_success(token)
|
| else:
|
| await self.record_token_failure(
|
| token,
|
| Exception(error_message or "验证失败"),
|
| )
|
|
|
| return is_valid
|
|
|
| async def health_check_all(self):
|
| """异步健康检查所有 Token"""
|
| if not self.token_statuses:
|
| logger.debug("⚠️ Token 池为空，跳过健康检查")
|
| return
|
|
|
| total_tokens = len(self.token_statuses)
|
| concurrency = 10
|
| started_at = time.perf_counter()
|
| logger.info(
|
| "[token.health_check] start total_tokens=%s mode=concurrency concurrency=%s",
|
| total_tokens,
|
| concurrency,
|
| )
|
|
|
|
|
| semaphore = asyncio.Semaphore(concurrency)
|
|
|
| async def _check_with_limit(token: str):
|
| async with semaphore:
|
| return await self.health_check_token(token)
|
|
|
| tasks = [
|
| _check_with_limit(token)
|
| for token in self.token_statuses.keys()
|
| ]
|
|
|
| results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
|
| healthy_count = sum(1 for r in results if r is True)
|
| failed_count = sum(1 for r in results if r is False)
|
| exception_count = sum(1 for r in results if isinstance(r, Exception))
|
|
|
| health_rate = (healthy_count / total_tokens) * 100 if total_tokens > 0 else 0
|
|
|
| elapsed_ms = (time.perf_counter() - started_at) * 1000
|
| logger.info(
|
| "[token.health_check] done total_tokens=%s healthy=%s failed=%s exceptions=%s mode=concurrency concurrency=%s elapsed_ms=%.1f",
|
| total_tokens,
|
| healthy_count,
|
| failed_count,
|
| exception_count,
|
| concurrency,
|
| elapsed_ms,
|
| )
|
|
|
| if healthy_count == 0 and total_tokens > 0:
|
| logger.warning(f"⚠️ 健康检查完成: 0/{total_tokens} 个 Token 健康 - 请检查 Token 配置")
|
| elif failed_count > 0:
|
| logger.warning(f"⚠️ 健康检查完成: {healthy_count}/{total_tokens} 个 Token 健康 ({health_rate:.1f}%)")
|
| else:
|
| logger.info(f"✅ 健康检查完成: {healthy_count}/{total_tokens} 个 Token 健康")
|
|
|
| if exception_count > 0:
|
| logger.error(f"💥 {exception_count} 个 Token 检查异常")
|
|
|
| async def sync_from_database(self, provider: str = "zai"):
|
| """
|
| 从数据库同步 Token 状态（禁用/启用状态）
|
|
|
| Args:
|
| provider: 提供商名称
|
|
|
| 说明：
|
| - 从数据库读取最新的 Token 启用状态
|
| - 如果数据库中 Token 被禁用，则从池中移除
|
| - 如果数据库中有新增的启用 Token，则添加到池中
|
| - 保留现有 Token 的运行时统计（请求数、成功率等）
|
| """
|
| from app.services.token_dao import get_token_dao
|
|
|
| dao = get_token_dao()
|
|
|
|
|
| token_records = await dao.get_tokens_by_provider(provider, enabled_only=True)
|
|
|
|
|
| db_tokens = {
|
| record["token"]: (
|
| record["id"],
|
| record.get("token_type", "unknown"),
|
| int(record.get("priority", 0) or 0),
|
| )
|
| for record in token_records
|
| if record.get("token_type") != "guest"
|
| }
|
|
|
| async with self._lock:
|
|
|
| should_reset_weights = False
|
| tokens_to_remove = []
|
| for token_value in list(self.token_statuses.keys()):
|
| if token_value not in db_tokens:
|
| tokens_to_remove.append(token_value)
|
|
|
| for token_value in tokens_to_remove:
|
| del self.token_statuses[token_value]
|
| del self.token_id_map[token_value]
|
| if tokens_to_remove:
|
| should_reset_weights = True
|
|
|
|
|
| new_tokens_count = 0
|
| ordered_user_tokens: List[str] = []
|
| for token_value, (token_id, token_type, priority) in db_tokens.items():
|
| if token_type == "user":
|
| ordered_user_tokens.append(token_value)
|
| if token_value not in self.token_statuses:
|
| self.token_statuses[token_value] = TokenStatus(
|
| token=token_value,
|
| token_id=token_id,
|
| token_type=token_type,
|
| priority=priority,
|
| )
|
| self.token_id_map[token_value] = token_id
|
| new_tokens_count += 1
|
| should_reset_weights = True
|
|
|
|
|
| for token_value, (token_id, token_type, priority) in db_tokens.items():
|
| if token_value in self.token_statuses:
|
| status = self.token_statuses[token_value]
|
| old_type = status.token_type
|
| if old_type != token_type:
|
| status.token_type = token_type
|
| should_reset_weights = True
|
| if status.priority != priority:
|
| status.priority = priority
|
| should_reset_weights = True
|
|
|
| if should_reset_weights and self.strategy == TOKEN_LOAD_BALANCE_SMOOTH_WEIGHTED:
|
| for status in self.token_statuses.values():
|
| status.current_weight = 0
|
|
|
| self._user_tokens = ordered_user_tokens
|
| if self._user_tokens:
|
| self._current_index %= len(self._user_tokens)
|
| else:
|
| self._current_index = 0
|
|
|
| logger.info(
|
| f"✅ Token 池同步完成: "
|
| f"当前 {len(self.token_statuses)} 个 Token "
|
| f"(移除 {len(tokens_to_remove)}, 新增 {new_tokens_count})"
|
| )
|
|
|
|
|
|
|
|
|
|
|
| _token_pool: Optional[TokenPool] = None
|
| _pool_lock = Lock()
|
|
|
|
|
| def get_token_pool() -> Optional[TokenPool]:
|
| """获取全局 Token 池实例"""
|
| return _token_pool
|
|
|
|
|
| async def initialize_token_pool_from_db(
|
| provider: str = "zai",
|
| failure_threshold: int = 3,
|
| recovery_timeout: int = 1800,
|
| strategy: str = TOKEN_LOAD_BALANCE_SMOOTH_WEIGHTED,
|
| ) -> Optional[TokenPool]:
|
| """
|
| 从数据库初始化全局 Token 池
|
|
|
| Args:
|
| provider: 提供商名称（当前仅使用 zai）
|
| failure_threshold: 失败阈值
|
| recovery_timeout: 恢复超时时间（秒）
|
|
|
| Returns:
|
| TokenPool 实例（即使没有 Token 也会创建空池）
|
| """
|
| global _token_pool
|
|
|
| from app.services.token_dao import get_token_dao
|
|
|
| dao = get_token_dao()
|
|
|
|
|
| token_records = await dao.get_tokens_by_provider(provider, enabled_only=True)
|
|
|
|
|
| tokens = []
|
| if token_records:
|
| tokens = [
|
| (
|
| record["id"],
|
| record["token"],
|
| record.get("token_type", "unknown"),
|
| int(record.get("priority", 0) or 0),
|
| )
|
| for record in token_records
|
| ]
|
|
|
|
|
| user_tokens = [
|
| token_record
|
| for token_record in tokens
|
| if str(token_record[2]) != "guest"
|
| ]
|
|
|
| if len(user_tokens) < len(tokens):
|
| guest_count = len(tokens) - len(user_tokens)
|
| logger.debug("过滤了 %s 个匿名用户 Token", guest_count)
|
|
|
| tokens = user_tokens
|
|
|
|
|
| async with _pool_lock:
|
| _token_pool = TokenPool(
|
| tokens,
|
| failure_threshold,
|
| recovery_timeout,
|
| strategy=strategy,
|
| )
|
|
|
| if not tokens:
|
| logger.warning(f"⚠️ {provider} 没有有效的认证用户 Token，已创建空 Token 池")
|
| else:
|
| logger.debug("从数据库初始化 Token 池（%s），共 %s 个 Token", provider, len(tokens))
|
|
|
| return _token_pool
|
|
|
|
|
| async def sync_token_stats_to_db():
|
| """
|
| 将内存中的 Token 统计同步到数据库
|
|
|
| 应在服务关闭或定期调用，确保统计数据不丢失
|
| """
|
| pool = get_token_pool()
|
| if not pool:
|
| return
|
|
|
| from app.services.token_stats_writer import get_token_stats_writer
|
|
|
| await get_token_stats_writer().flush_now()
|
|
|
| from app.services.token_dao import get_token_dao
|
|
|
| dao = get_token_dao()
|
|
|
| pending_updates = []
|
| async with pool._lock:
|
| for token, status in pool.token_statuses.items():
|
| pending_success = max(
|
| 0,
|
| status.successful_requests - status.db_synced_successful_requests,
|
| )
|
| pending_failure = max(
|
| 0,
|
| status.failed_requests - status.db_synced_failed_requests,
|
| )
|
| if pending_success > 0 or pending_failure > 0:
|
| pending_updates.append(
|
| (
|
| token,
|
| status.token_id,
|
| pending_success,
|
| pending_failure,
|
| )
|
| )
|
|
|
| for token, token_id, pending_success, pending_failure in pending_updates:
|
| for _ in range(pending_success):
|
| await dao.record_success(token_id)
|
| for _ in range(pending_failure):
|
| await dao.record_failure(token_id)
|
|
|
| async with pool._lock:
|
| if token in pool.token_statuses:
|
| status = pool.token_statuses[token]
|
| status.db_synced_successful_requests += pending_success
|
| status.db_synced_failed_requests += pending_failure
|
|
|
| logger.info("✅ Token 统计已同步到数据库")
|
|
|
