Spaces:

bg5
/

z

Sleeping

z / app /utils /token_pool.py
Upload 85 files
712a3c7 verified 6 months ago
36.6 kB
#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
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 状态管理 ====================


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  # 数据库 ID，用于同步统计
    token_type: str = "unknown"  # "user", "guest", "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 永远不健康
        """
        # guest 和 unknown token 永远不健康
        if self.token_type != "user":
            return False

        # 不可用的 token 不健康
        if not self.is_available:
            return False

        # 新 token 容错：请求数很少时，只要没触发连续错误熔断就健康
        if self.total_requests <= 3:
            return self.failure_count == 0

        # 基于成功率判断
        return self.success_rate >= 0.5


# ==================== Token 验证服务 ====================


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
        """
        # 检查 HTTP 状态码
        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 字段
            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)}")


# ==================== Token 池管理器 ====================


class TokenPool:
    """Token 池管理器（数据库驱动）"""

    def __init__(
        self,
        tokens: List[Tuple[int, str, str]],  # [(token_id, token_value, token_type[, priority]), ...]
        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)

        # 初始化 Token 状态（内存中）
        self.token_statuses: Dict[str, TokenStatus] = {}
        self.token_id_map: Dict[str, int] = {}  # token -> token_id 映射
        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)

            # 统计各类型 Token
            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)

        # 更新 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,
        )

        # 并发执行所有 Token 的健康检查
        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
        token_records = await dao.get_tokens_by_provider(provider, enabled_only=True)

        # 构建数据库中的 Token 映射
        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"  # 过滤 guest token
        }

        async with self._lock:
            # 1. 移除已在数据库中禁用的 Token
            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

            # 2. 添加新启用的 Token
            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

            # 3. 更新现有 Token 的类型（如果数据库中有更新）
            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（只加载启用的认证用户 Token）
    token_records = await dao.get_tokens_by_provider(provider, enabled_only=True)

    # 转换为 TokenPool 所需格式
    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
        ]

        # 过滤掉 guest token（不应该在数据库中，但防御性检查）
        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

    # 始终创建 Token 池实例（即使为空）
    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 统计已同步到数据库")