app/utils/token_pool.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Token池管理器
实现AUTH_TOKEN的轮询机制，提供负载均衡和容错功能
"""

import asyncio
import time
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field
from threading import Lock
import httpx

from app.utils.logger import logger


@dataclass
class TokenStatus:
    """Token状态信息"""
    token: str
    is_available: bool = True
    failure_count: int = 0
    last_failure_time: float = 0.0
    last_success_time: float = 0.0
    total_requests: int = 0
    successful_requests: int = 0
    token_type: str = "unknown"  # "user", "guest", "unknown"
    
    @property
    def success_rate(self) -> float:
        """成功率"""
        if self.total_requests == 0:
            return 1.0
        return self.successful_requests / self.total_requests
    
    @property
    def is_healthy(self) -> bool:
        """
        是否健康

        健康的定义：
        1. 必须是认证用户token (token_type = "user")
        2. 当前可用 (is_available = True)
        3. 成功率 >= 50% 或者总请求数 <= 3（新token容错）

        注意：guest token不应该在AUTH_TOKENS中
        """
        # guest token永远不健康
        if self.token_type == "guest":
            return False

        # 未知类型token不健康
        if self.token_type != "user":
            return False

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

        # 对于认证用户token，基于成功率判断
        # 新token或请求数很少时，给予容错
        if self.total_requests <= 3:
            return self.failure_count == 0

        # 基于成功率判断健康状态
        return self.success_rate >= 0.5


class TokenPool:
    """Token池管理器"""
    
    def __init__(self, tokens: List[str], failure_threshold: int = 3, recovery_timeout: int = 1800):
        """
        初始化Token池
        
        Args:
            tokens: token列表
            failure_threshold: 失败阈值，超过此次数将标记为不可用
            recovery_timeout: 恢复超时时间（秒），失败token在此时间后重新尝试
        """
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self._lock = Lock()
        self._current_index = 0
        
        # 初始化token状态
        self.token_statuses: Dict[str, TokenStatus] = {}
        original_count = len(tokens)
        unique_tokens = []

        # 去重处理
        for token in tokens:
            if token and token not in self.token_statuses:  # 过滤空token和重复token
                # 预设为认证用户token，因为这些是用户手动配置的token
                self.token_statuses[token] = TokenStatus(token=token, token_type="user")
                unique_tokens.append(token)

        duplicate_count = original_count - len(unique_tokens)
        if duplicate_count > 0:
            logger.warning(f"⚠️ 检测到 {duplicate_count} 个重复token，已自动去重")

        if not self.token_statuses:
            logger.warning("⚠️ Token池为空，将依赖匿名模式")
        # else:
        #     logger.info(f"🔧 初始化Token池，共 {len(self.token_statuses)} 个token")
    
    def get_next_token(self) -> Optional[str]:
        """
        获取下一个可用的token（轮询算法）

        Returns:
            可用的token，如果没有可用token则返回None
        """
        with self._lock:
            if not self.token_statuses:
                return None

            available_tokens = self._get_available_tokens()
            if not available_tokens:
                # 尝试恢复过期的失败token
                self._try_recover_failed_tokens()
                available_tokens = self._get_available_tokens()

                if not available_tokens:
                    logger.warning("⚠️ 没有可用的token")
                    return None

            # 轮询选择token
            token = available_tokens[self._current_index % len(available_tokens)]
            self._current_index = (self._current_index + 1) % len(available_tokens)

            return token
    
    def _get_available_tokens(self) -> List[str]:
        """
        获取当前可用的认证用户token列表

        返回满足以下条件的token：
        1. is_available = True (可用状态)
        2. token_type == "user" (认证用户token)

        这确保轮询机制只会选择有效的认证用户token，跳过匿名用户token
        """
        available_user_tokens = [
            status.token for status in self.token_statuses.values()
            if status.is_available and status.token_type == "user"
        ]

        # 检查是否有匿名用户token并给出警告
        if not available_user_tokens and self.token_statuses:
            guest_tokens = [
                status.token for status in self.token_statuses.values()
                if status.token_type == "guest"
            ]
            if guest_tokens:
                logger.warning(f"⚠️ 检测到 {len(guest_tokens)} 个匿名用户token，轮询机制将跳过这些token")

        return available_user_tokens
    
    def _try_recover_failed_tokens(self):
        """尝试恢复失败的token"""
        current_time = time.time()
        recovered_count = 0
        
        for status in self.token_statuses.values():
            if (not status.is_available and 
                current_time - status.last_failure_time > self.recovery_timeout):
                status.is_available = True
                status.failure_count = 0
                recovered_count += 1
                logger.info(f"🔄 恢复失败token: {status.token[:20]}...")
        
        if recovered_count > 0:
            logger.info(f"✅ 恢复了 {recovered_count} 个失败的token")
    
    def mark_token_success(self, token: str):
        """标记token使用成功"""
        with self._lock:
            if token in self.token_statuses:
                status = self.token_statuses[token]
                status.total_requests += 1
                status.successful_requests += 1
                status.last_success_time = time.time()
                status.failure_count = 0  # 重置失败计数
                
                if not status.is_available:
                    status.is_available = True
                    logger.info(f"✅ Token恢复可用: {token[:20]}...")
    
    def mark_token_failure(self, token: str, error: Exception = None):
        """标记token使用失败"""
        with self._lock:
            if token in self.token_statuses:
                status = self.token_statuses[token]
                status.total_requests += 1
                status.failure_count += 1
                status.last_failure_time = time.time()
                
                if status.failure_count >= self.failure_threshold:
                    status.is_available = False
                    logger.warning(f"🚫 Token已禁用: {token[:20]}... (失败 {status.failure_count} 次)")
    
    def get_pool_status(self) -> Dict:
        """获取token池状态信息"""
        with self._lock:
            available_count = len(self._get_available_tokens())
            total_count = len(self.token_statuses)

            # 统计健康token数量
            healthy_count = sum(1 for status in self.token_statuses.values() if status.is_healthy)

            status_info = {
                "total_tokens": total_count,
                "available_tokens": available_count,
                "unavailable_tokens": total_count - available_count,
                "healthy_tokens": healthy_count,
                "unhealthy_tokens": total_count - healthy_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_type": status.token_type,
                    "is_available": status.is_available,
                    "failure_count": status.failure_count,
                    "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
                })

            return status_info
    
    def update_tokens(self, new_tokens: List[str]):
        """动态更新token列表"""
        with self._lock:
            # 保留现有token的状态信息
            old_statuses = self.token_statuses.copy()
            self.token_statuses.clear()

            original_count = len(new_tokens)
            unique_tokens = []

            # 去重并添加新token，保留已存在token的状态
            for token in new_tokens:
                if token and token not in self.token_statuses:  # 过滤空token和重复token
                    if token in old_statuses:
                        self.token_statuses[token] = old_statuses[token]
                    else:
                        # 预设为认证用户token，因为这些是用户手动配置的token
                        self.token_statuses[token] = TokenStatus(token=token, token_type="user")
                    unique_tokens.append(token)

            # 记录去重信息
            duplicate_count = original_count - len(unique_tokens)
            if duplicate_count > 0:
                logger.warning(f"⚠️ 更新时检测到 {duplicate_count} 个重复token，已自动去重")

            # 重置索引
            self._current_index = 0

            logger.info(f"🔄 更新Token池，共 {len(self.token_statuses)} 个token")
    
    async def health_check_token(self, token: str, auth_url: str = "https://chat.z.ai/api/v1/auths/") -> bool:
        """
        异步健康检查单个token

        使用Z.AI认证API验证token的有效性，通过检查响应内容判断token是否有效

        Args:
            token: 要检查的token
            auth_url: 认证URL

        Returns:
            token是否健康
        """
        try:
            # 构建完整的请求头，模拟真实浏览器请求
            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"
            }

            async with httpx.AsyncClient(timeout=15.0) as client:
                response = await client.get(auth_url, headers=headers)

                # 验证token有效性并获取类型
                token_type, is_healthy = self._validate_token_response(response)

                # 更新token类型
                if token in self.token_statuses:
                    self.token_statuses[token].token_type = token_type

                if is_healthy:
                    self.mark_token_success(token)
                else:
                    # 简化错误信息，只记录关键错误类型
                    if token_type == "guest":
                        error_msg = "匿名用户token"
                    elif response.status_code != 200:
                        error_msg = f"HTTP {response.status_code}"
                    else:
                        error_msg = "认证失败"

                    self.mark_token_failure(token, Exception(error_msg))

                return is_healthy

        except (httpx.TimeoutException, httpx.ConnectError, Exception) as e:
            self.mark_token_failure(token, e)
            return False

    def _validate_token_response(self, response: httpx.Response) -> bool:
        """
        基于Z.AI API响应中的role字段验证token类型

        验证规则：
        - role: "user" = 认证用户token（有效，可用于AUTH_TOKENS）
        - role: "guest" = 匿名用户token（无效，不应在AUTH_TOKENS中）
        - 无role字段或其他值 = 无效token

        Args:
            response: HTTP响应对象

        Returns:
            token是否为有效的认证用户token
        """
        # 首先检查HTTP状态码
        if response.status_code != 200:
            return ("unknown", False)

        try:
            # 尝试解析JSON响应
            response_data = response.json()

            if not isinstance(response_data, dict):
                return ("unknown", False)

            # 检查是否包含错误信息
            if "error" in response_data:
                return ("unknown", False)

            if "message" in response_data and "error" in response_data.get("message", "").lower():
                return ("unknown", False)

            # 核心验证：检查role字段
            role = response_data.get("role")

            if role == "user":
                return ("user", True)
            elif role == "guest":
            
                if not hasattr(self, '_guest_token_warned'):
                    logger.warning("⚠️ 检测到匿名用户token，建议仅在AUTH_TOKENS中配置认证用户token")
                    self._guest_token_warned = True
                return ("guest", False)
            else:
                return ("unknown", False)

        except (ValueError, Exception):
            return ("unknown", False)

    async def health_check_all(self, auth_url: str = "https://chat.z.ai/api/v1/auths/"):
        """异步健康检查所有token"""
        if not self.token_statuses:
            logger.warning("⚠️ Token池为空，跳过健康检查")
            return

        total_tokens = len(self.token_statuses)
        logger.info(f"🔍 开始Token池健康检查... (共 {total_tokens} 个token)")

        # 并发执行所有token的健康检查
        tasks = []
        token_list = list(self.token_statuses.keys())

        for token in token_list:
            task = self.health_check_token(token, auth_url)
            tasks.append(task)

        # 执行并收集结果
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # 统计结果
        healthy_count = 0
        failed_count = 0
        exception_count = 0

        for i, result in enumerate(results):
            if result is True:
                healthy_count += 1
            elif result is False:
                failed_count += 1
            else:
                # 异常情况
                exception_count += 1
                token = token_list[i]
                logger.error(f"💥 Token {token[:20]}... 健康检查异常: {result}")

        health_rate = (healthy_count / total_tokens) * 100 if total_tokens > 0 else 0

        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检查异常")


# 全局token池实例
_token_pool: Optional[TokenPool] = None
_pool_lock = Lock()


def get_token_pool() -> Optional[TokenPool]:
    """获取全局token池实例"""
    return _token_pool


def initialize_token_pool(tokens: List[str], failure_threshold: int = 3, recovery_timeout: int = 1800) -> TokenPool:
    """初始化全局token池"""
    global _token_pool
    with _pool_lock:
        _token_pool = TokenPool(tokens, failure_threshold, recovery_timeout)
        return _token_pool


def update_token_pool(tokens: List[str]):
    """更新全局token池"""
    global _token_pool
    with _pool_lock:
        if _token_pool:
            _token_pool.update_tokens(tokens)
        else:
            _token_pool = TokenPool(tokens)