gemini-balance/app/service/key/key_manager.py

import asyncio
from itertools import cycle
from typing import Dict

from app.config.config import settings
from app.log.logger import get_key_manager_logger

logger = get_key_manager_logger()


class KeyManager:
    def __init__(self, api_keys: list):
        self.api_keys = api_keys
        self.key_cycle = cycle(api_keys)
        self.key_cycle_lock = asyncio.Lock()
        self.failure_count_lock = asyncio.Lock()
        self.key_failure_counts: Dict[str, int] = {key: 0 for key in api_keys}
        self.MAX_FAILURES = settings.MAX_FAILURES
        self.paid_key = settings.PAID_KEY

    async def get_paid_key(self) -> str:
        return self.paid_key

    async def get_next_key(self) -> str:
        """获取下一个API key"""
        async with self.key_cycle_lock:
            return next(self.key_cycle)

    async def is_key_valid(self, key: str) -> bool:
        """检查key是否有效"""
        async with self.failure_count_lock:
            return self.key_failure_counts[key] < self.MAX_FAILURES

    async def reset_failure_counts(self):
        """重置所有key的失败计数"""
        async with self.failure_count_lock:
            for key in self.key_failure_counts:
                self.key_failure_counts[key] = 0

    async def reset_key_failure_count(self, key: str) -> bool:
        """重置指定key的失败计数"""
        async with self.failure_count_lock:
            if key in self.key_failure_counts:
                self.key_failure_counts[key] = 0
                logger.info(f"Reset failure count for key: {key}")
                return True
            logger.warning(
                f"Attempt to reset failure count for non-existent key: {key}"
            )
            return False

    async def get_next_working_key(self) -> str:
        """获取下一可用的API key"""
        initial_key = await self.get_next_key()
        current_key = initial_key

        while True:
            if await self.is_key_valid(current_key):
                return current_key

            current_key = await self.get_next_key()
            if current_key == initial_key:
                # await self.reset_failure_counts() 取消重置
                return current_key

    async def handle_api_failure(self, api_key: str, retries: int) -> str:
        """处理API调用失败"""
        async with self.failure_count_lock:
            self.key_failure_counts[api_key] += 1
            if self.key_failure_counts[api_key] >= self.MAX_FAILURES:
                logger.warning(
                    f"API key {api_key} has failed {self.MAX_FAILURES} times"
                )
        if retries < settings.MAX_RETRIES:
            return await self.get_next_working_key()
        else:
            return ""

    def get_fail_count(self, key: str) -> int:
        """获取指定密钥的失败次数"""
        return self.key_failure_counts.get(key, 0)

    async def get_keys_by_status(self) -> dict:
        """获取分类后的API key列表，包括失败次数"""
        valid_keys = {}
        invalid_keys = {}

        async with self.failure_count_lock:
            for key in self.api_keys:
                fail_count = self.key_failure_counts[key]
                if fail_count < self.MAX_FAILURES:
                    valid_keys[key] = fail_count
                else:
                    invalid_keys[key] = fail_count

        return {"valid_keys": valid_keys, "invalid_keys": invalid_keys}

    async def get_first_valid_key(self) -> str:
        """获取第一个有效的API key"""
        async with self.failure_count_lock:
            for key in self.key_failure_counts:
                if self.key_failure_counts[key] < self.MAX_FAILURES:
                    return key
        # 如果所有 key 都无效，或者列表为空，则尝试返回第一个（如果列表不为空）
        # 或者根据具体逻辑处理，这里保持原样，可能在空列表或全无效时需要调整
        if self.api_keys:
            return self.api_keys[0]
        # 如果 api_keys 为空，这里会出问题。实际应用中应有非空保证或更好处理。
        # 为了保持接口一致性，如果列表为空，可能应该抛出异常或返回特定值。
        # 暂且假设 api_keys 不会为空，或者调用者处理后续的空 key 问题。
        # 根据现有代码，如果api_keys为空，self.api_keys[0]会报错。
        # 如果没有有效key且列表不空，返回第一个。若列表为空，这里会出IndexError。
        # 更安全的做法是：
        if not self.api_keys:
            logger.warning("API key list is empty, cannot get first valid key.")
            # Depending on desired behavior, either raise error or return an indicator like "" or None
            # For now, let's allow it to potentially fail if a key is expected by caller
            # but it's better to be explicit. Let's return empty string for consistency with handle_api_failure
            return ""
        return self.api_keys[
            0
        ]  # Fallback to the first key if no key is "valid" but list is not empty


_singleton_instance = None
_singleton_lock = asyncio.Lock()
_preserved_failure_counts: Dict[str, int] | None = None
_preserved_old_api_keys_for_reset: list | None = None
_preserved_next_key_in_cycle: str | None = None


async def get_key_manager_instance(api_keys: list = None) -> KeyManager:
    """
    获取 KeyManager 单例实例。

    如果尚未创建实例，将使用提供的 api_keys 初始化 KeyManager。
    如果已创建实例，则忽略 api_keys 参数，返回现有单例。
    如果在重置后调用，会尝试恢复之前的状态（失败计数、循环位置）。
    """
    global _singleton_instance, _preserved_failure_counts, _preserved_old_api_keys_for_reset, _preserved_next_key_in_cycle

    async with _singleton_lock:
        if _singleton_instance is None:
            if api_keys is None:
                # This case needs careful handling. If it's the very first call, api_keys are required.
                # If it's after a reset and no api_keys are provided, what should happen?
                # The original ValueError was "API keys are required to initialize the KeyManager".
                # Let's assume if api_keys is None here, it's an error unless we are restoring from non-None _preserved_old_api_keys_for_reset.
                # However, the user's request implies new api_keys will be part of the reset flow.
                # For now, stick to a strict requirement for api_keys if _singleton_instance is None.
                raise ValueError(
                    "API keys are required to initialize or re-initialize the KeyManager instance."
                )
            if not api_keys:  # Handle case where api_keys is an empty list
                logger.warning(
                    "Initializing KeyManager with an empty list of API keys."
                )
                # Consider if this should be an error or allowed. Current KeyManager supports it.

            _singleton_instance = KeyManager(api_keys)
            logger.info(
                f"KeyManager instance created/re-created with {len(api_keys)} API keys."
            )

            # 1. 恢复失败计数
            if _preserved_failure_counts:
                # Initialize new instance's failure_counts for all new keys to 0
                current_failure_counts = {
                    key: 0 for key in _singleton_instance.api_keys
                }
                # Inherit counts for keys that exist in both old and new lists
                for key, count in _preserved_failure_counts.items():
                    if key in current_failure_counts:
                        current_failure_counts[key] = count
                _singleton_instance.key_failure_counts = current_failure_counts
                logger.info("Inherited failure counts for applicable keys.")
            _preserved_failure_counts = None  # Clear after use

            # 2. 调整 key_cycle 的起始点
            start_key_for_new_cycle = None
            if (
                _preserved_old_api_keys_for_reset
                and _preserved_next_key_in_cycle
                and _singleton_instance.api_keys  # Ensure new api_keys list is not empty
            ):
                try:
                    # Find the index of the preserved next key in the *old* list
                    start_idx_in_old = _preserved_old_api_keys_for_reset.index(
                        _preserved_next_key_in_cycle
                    )

                    # Iterate through the old key list (circularly) starting from _preserved_next_key_in_cycle
                    # Find the first key that also exists in the new api_keys list
                    for i in range(len(_preserved_old_api_keys_for_reset)):
                        current_old_key_idx = (start_idx_in_old + i) % len(
                            _preserved_old_api_keys_for_reset
                        )
                        key_candidate = _preserved_old_api_keys_for_reset[
                            current_old_key_idx
                        ]
                        if key_candidate in _singleton_instance.api_keys:
                            start_key_for_new_cycle = key_candidate
                            break
                except ValueError:
                    logger.warning(
                        f"Preserved next key '{_preserved_next_key_in_cycle}' not found in preserved old API keys. "
                        "New cycle will start from the beginning of the new list."
                    )
                except Exception as e:
                    logger.error(
                        f"Error determining start key for new cycle from preserved state: {e}. "
                        "New cycle will start from the beginning."
                    )

            if start_key_for_new_cycle and _singleton_instance.api_keys:
                try:
                    # Find the index of the determined start_key in the new api_keys list
                    target_idx = _singleton_instance.api_keys.index(
                        start_key_for_new_cycle
                    )
                    # Advance the new cycle by calling next() target_idx times
                    # This positions the cycle so that the *next* call to next() will yield start_key_for_new_cycle
                    for _ in range(target_idx):
                        next(_singleton_instance.key_cycle)
                    logger.info(
                        f"Key cycle in new instance advanced. Next call to get_next_key() will yield: {start_key_for_new_cycle}"
                    )
                except ValueError:
                    # This should not happen if start_key_for_new_cycle was correctly found in api_keys
                    logger.warning(
                        f"Determined start key '{start_key_for_new_cycle}' not found in new API keys during cycle advancement. "
                        "New cycle will start from the beginning."
                    )
                except (
                    StopIteration
                ):  # Should not happen with cycle unless api_keys is empty, handled by _singleton_instance.api_keys check
                    logger.error(
                        "StopIteration while advancing key cycle, implies empty new API key list previously missed."
                    )
                except Exception as e:
                    logger.error(
                        f"Error advancing new key cycle: {e}. Cycle will start from beginning."
                    )
            else:
                if _singleton_instance.api_keys:
                    logger.info(
                        "New key cycle will start from the beginning of the new API key list (no specific start key determined or needed)."
                    )
                else:
                    logger.info(
                        "New key cycle not applicable as the new API key list is empty."
                    )

            # 清理所有保存的状态
            _preserved_old_api_keys_for_reset = None
            _preserved_next_key_in_cycle = None
            # _preserved_failure_counts already cleared

        return _singleton_instance


async def reset_key_manager_instance():
    """
    重置 KeyManager 单例实例。
    将保存当前实例的状态（失败计数、旧 API keys、下一个 key 提示）
    以供下一次 get_key_manager_instance 调用时恢复。
    """
    global _singleton_instance, _preserved_failure_counts, _preserved_old_api_keys_for_reset, _preserved_next_key_in_cycle
    async with _singleton_lock:
        if _singleton_instance:
            # 1. 保存失败计数
            _preserved_failure_counts = _singleton_instance.key_failure_counts.copy()

            # 2. 保存旧的 API keys 列表
            _preserved_old_api_keys_for_reset = _singleton_instance.api_keys.copy()

            # 3. 保存 key_cycle 的下一个 key 提示
            # This should be the key that get_next_key() would return next.
            try:
                if (
                    _singleton_instance.api_keys
                ):  # Only if there are keys to cycle through
                    # Calling get_next_key() consumes one key and returns it. This is the key
                    # we want the new cycle to effectively start with.
                    _preserved_next_key_in_cycle = (
                        await _singleton_instance.get_next_key()
                    )
                else:
                    _preserved_next_key_in_cycle = None  # No keys, so no next key
            except (
                StopIteration
            ):  # Should be caught by "if _singleton_instance.api_keys"
                logger.warning(
                    "Could not preserve next key hint: key cycle was empty or exhausted in old instance."
                )
                _preserved_next_key_in_cycle = None
            except Exception as e:
                logger.error(f"Error preserving next key hint during reset: {e}")
                _preserved_next_key_in_cycle = None

            _singleton_instance = None
            logger.info(
                "KeyManager instance has been reset. State (failure counts, old keys, next key hint) preserved for next instantiation."
            )
        else:
            logger.info(
                "KeyManager instance was not set (or already reset), no reset action performed."
            )