""" MongoDB 存储管理器 """ import json import os import random import time from typing import Any, Dict, List, Optional from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase from log import log class MongoDBManager: """MongoDB 数据库管理器""" # 状态字段常量 STATE_FIELDS = { "error_codes", "error_messages", "disabled", "last_success", "user_email", "model_cooldowns", "preview", "tier", "enable_credit", } @staticmethod def _escape_model_name(model_name: str) -> str: """ 转义模型名中的点号,避免 MongoDB 将其解释为嵌套结构 Args: model_name: 原始模型名 (如 "gemini-2.5-flash") Returns: 转义后的模型名 (如 "gemini-2-5-flash") """ return model_name.replace(".", "-") def __init__(self): self._client: Optional[AsyncIOMotorClient] = None self._db: Optional[AsyncIOMotorDatabase] = None self._initialized = False # 内存配置缓存 - 初始化时加载一次 self._config_cache: Dict[str, Any] = {} self._config_loaded = False # Redis 缓存（仅当 REDIS_URL 环境变量存在时启用） self._redis = None self._redis_enabled: bool = False async def initialize(self) -> None: """初始化 MongoDB 连接""" if self._initialized: return try: mongodb_uri = os.getenv("MONGODB_URI") if not mongodb_uri: raise ValueError("MONGODB_URI environment variable not set") database_name = os.getenv("MONGODB_DATABASE", "gcli2api") self._client = AsyncIOMotorClient(mongodb_uri) self._db = self._client[database_name] # 测试连接 await self._db.command("ping") # 创建索引 await self._create_indexes() # 加载配置到内存 await self._load_config_cache() self._initialized = True log.info(f"MongoDB storage initialized (database: {database_name})") # 尝试初始化 Redis（可选） await self._init_redis() except Exception as e: log.error(f"Error initializing MongoDB: {e}") raise async def _create_indexes(self): """ 创建索引 """ from pymongo import IndexModel, ASCENDING credentials_collection = self._db["credentials"] antigravity_credentials_collection = self._db["antigravity_credentials"] # ===== Geminicli 凭证索引 ===== geminicli_indexes = [ # 唯一索引 - 用于所有按文件名的精确查询 IndexModel([("filename", ASCENDING)], unique=True, name="idx_filename_unique"), # 复合索引 - 用于 get_next_available_credential 和 get_available_credentials_list # 查询模式: {disabled: False} + sort by rotation_order IndexModel( [("disabled", ASCENDING), ("rotation_order", ASCENDING)], name="idx_disabled_rotation" ), # 单字段索引 - 用于 get_credentials_summary 的错误筛选 IndexModel([("error_codes", ASCENDING)], name="idx_error_codes"), # 单字段索引 - 用于 get_duplicate_credentials_by_email 的去重查询 IndexModel([("user_email", ASCENDING)], name="idx_user_email"), ] # ===== Antigravity 凭证索引 ===== antigravity_indexes = [ # 唯一索引 IndexModel([("filename", ASCENDING)], unique=True, name="idx_filename_unique"), # 复合索引 - 查询模式: {disabled: False} + sort by rotation_order # 查询模式: {disabled: False} + 可选 sort by rotation_order IndexModel( [("disabled", ASCENDING), ("rotation_order", ASCENDING)], name="idx_disabled_rotation" ), # 单字段索引 - 错误筛选 IndexModel([("error_codes", ASCENDING)], name="idx_error_codes"), # 单字段索引 - 去重查询 IndexModel([("user_email", ASCENDING)], name="idx_user_email"), ] # 并行创建新索引 try: await credentials_collection.create_indexes(geminicli_indexes) await antigravity_credentials_collection.create_indexes(antigravity_indexes) log.debug("MongoDB indexes created successfully") except Exception as e: # 如果索引已存在，忽略错误 if "already exists" not in str(e).lower(): log.warning(f"Index creation warning: {e}") async def _load_config_cache(self): """加载配置到内存缓存（仅在初始化时调用一次）""" if self._config_loaded: return try: config_collection = self._db["config"] cursor = config_collection.find({}) async for doc in cursor: self._config_cache[doc["key"]] = doc.get("value") self._config_loaded = True log.debug(f"Loaded {len(self._config_cache)} config items into cache") except Exception as e: log.error(f"Error loading config cache: {e}") self._config_cache = {} # ============ Redis 缓存（可选，仅当 REDIS_URL 存在时启用）============ async def _init_redis(self) -> None: """初始化 Redis 连接并重建凭证池缓存（若 REDIS_URL 存在）""" redis_url = os.getenv("REDIS_URL") if not redis_url: return try: import redis.asyncio as aioredis # type: ignore except ImportError: log.warning("redis package not installed, Redis cache disabled. Run: pip install redis") return try: self._redis = aioredis.from_url(redis_url, decode_responses=True) await self._redis.ping() self._redis_enabled = True log.info("Redis connected, rebuilding credential pool cache...") # 并行重建两个 mode 的缓存及配置缓存 import asyncio await asyncio.gather( self._rebuild_redis_cache("geminicli"), self._rebuild_redis_cache("antigravity"), self._load_config_to_redis(), ) log.info("Redis credential pool cache ready") except Exception as e: log.warning(f"Redis init failed, falling back to MongoDB-only mode: {e}") self._redis = None self._redis_enabled = False # ---- Redis key 工具 ---- def _rk_avail(self, mode: str) -> str: """所有未禁用凭证的 Redis Set key""" return f"gcli:avail:{mode}" def _rk_tier(self, mode: str, tier: str) -> str: """按 tier 分桶的未禁用凭证 Redis Set key""" return f"gcli:tier:{mode}:{tier}" def _rk_preview(self, mode: str) -> str: """preview=True 凭证的 Redis Set key""" return f"gcli:preview:{mode}" def _rk_cd(self, mode: str, filename: str, escaped_model: str) -> str: """模型冷却 Redis key（带 TTL）""" return f"gcli:cd:{mode}:{filename}:{escaped_model}" # ---- Redis 缓存维护 ---- async def _rebuild_redis_cache(self, mode: str) -> None: """ 从 MongoDB 重建指定 mode 的 Redis 凭证池缓存。 使用临时 key + RENAME 原子替换 """ if not self._redis: return try: collection = self._db[self._get_collection_name(mode)] # 同时投影 model_cooldowns、tier、preview，以便重建缓存 projection: Dict[str, Any] = {"filename": 1, "disabled": 1, "model_cooldowns": 1, "tier": 1, "preview": 1, "_id": 0} avail: List[str] = [] tier_buckets: Dict[str, List[str]] = {} # tier -> [filename, ...] preview_members: List[str] = [] cooldown_entries: List[tuple] = [] # (cd_key, ttl_seconds, value) current_time = time.time() async for doc in collection.find({}, projection=projection): if not doc.get("disabled", False): filename = doc["filename"] avail.append(filename) # 按 tier 分桶 tier = doc.get("tier") or "pro" tier_buckets.setdefault(tier, []).append(filename) # preview 分桶（仅 geminicli） if mode == "geminicli" and doc.get("preview", True): preview_members.append(filename) # 收集未过期的模型冷却，重建 Redis TTL Key model_cooldowns = doc.get("model_cooldowns") or {} for escaped_model, cooldown_until in model_cooldowns.items(): if isinstance(cooldown_until, (int, float)) and cooldown_until > current_time: ttl = int(cooldown_until - current_time) if ttl > 0: cd_key = self._rk_cd(mode, filename, escaped_model) cooldown_entries.append((cd_key, ttl, str(cooldown_until))) tmp_avail = self._rk_avail(mode) + ":tmp" pipe = self._redis.pipeline() # 先写临时 key（此时正式 key 仍完整可用） pipe.delete(tmp_avail) if avail: pipe.sadd(tmp_avail, *avail) await pipe.execute() # RENAME 是原子操作：瞬间切换，不存在空窗 pipe2 = self._redis.pipeline() if avail: pipe2.rename(tmp_avail, self._rk_avail(mode)) else: pipe2.delete(self._rk_avail(mode)) pipe2.delete(tmp_avail) await pipe2.execute() # 重建 tier 分桶 Set（原子替换） all_tiers = ("free", "pro", "ultra") pipe3 = self._redis.pipeline() for tier in all_tiers: tier_key = self._rk_tier(mode, tier) tmp_tier_key = tier_key + ":tmp" pipe3.delete(tmp_tier_key) members = tier_buckets.get(tier, []) if members: pipe3.sadd(tmp_tier_key, *members) await pipe3.execute() pipe4 = self._redis.pipeline() for tier in all_tiers: tier_key = self._rk_tier(mode, tier) tmp_tier_key = tier_key + ":tmp" members = tier_buckets.get(tier, []) if members: pipe4.rename(tmp_tier_key, tier_key) else: pipe4.delete(tier_key) pipe4.delete(tmp_tier_key) await pipe4.execute() # 重建 preview 分桶（仅 geminicli） preview_key = self._rk_preview(mode) tmp_preview_key = preview_key + ":tmp" pipe5 = self._redis.pipeline() pipe5.delete(tmp_preview_key) if preview_members: pipe5.sadd(tmp_preview_key, *preview_members) await pipe5.execute() pipe6 = self._redis.pipeline() if preview_members: pipe6.rename(tmp_preview_key, preview_key) else: pipe6.delete(preview_key) pipe6.delete(tmp_preview_key) await pipe6.execute() # 批量恢复未过期的模型冷却 TTL Key if cooldown_entries: pipe7 = self._redis.pipeline() for cd_key, ttl, value in cooldown_entries: pipe7.setex(cd_key, ttl, value) await pipe7.execute() log.debug( f"Redis cache rebuilt [{mode}]: {len(avail)} avail, " f"tiers={{{', '.join(f'{t}:{len(tier_buckets.get(t, []))}' for t in all_tiers)}}}, " f"preview={len(preview_members)}, " f"{len(cooldown_entries)} cooldown key(s) restored" ) except Exception as e: log.warning(f"Redis rebuild cache error [{mode}]: {e}") async def _redis_add_cred(self, mode: str, filename: str, tier: str = "pro", preview: bool = True) -> None: """将凭证加入 Redis 可用池及对应 tier 分桶、preview 分桶""" if not self._redis_enabled: return try: pipe = self._redis.pipeline() pipe.sadd(self._rk_avail(mode), filename) pipe.sadd(self._rk_tier(mode, tier), filename) if mode == "geminicli" and preview: pipe.sadd(self._rk_preview(mode), filename) await pipe.execute() except Exception as e: log.warning(f"Redis add_cred error: {e}") async def _redis_remove_cred(self, mode: str, filename: str, tier: Optional[str] = None) -> None: """从 Redis 所有池中移除凭证""" if not self._redis_enabled: return try: pipe = self._redis.pipeline() pipe.srem(self._rk_avail(mode), filename) if tier: pipe.srem(self._rk_tier(mode, tier), filename) else: # tier 未知时从所有分桶中移除 for t in ("free", "pro", "ultra"): pipe.srem(self._rk_tier(mode, t), filename) pipe.srem(self._rk_preview(mode), filename) await pipe.execute() except Exception as e: log.warning(f"Redis remove_cred error: {e}") async def _redis_sync_cred(self, mode: str, filename: str, disabled: bool, tier: str = "pro", preview: bool = True) -> None: """根据最新状态同步单个凭证在 Redis 中的集合成员""" if not self._redis_enabled: return try: pipe = self._redis.pipeline() if disabled: pipe.srem(self._rk_avail(mode), filename) for t in ("free", "pro", "ultra"): pipe.srem(self._rk_tier(mode, t), filename) pipe.srem(self._rk_preview(mode), filename) else: pipe.sadd(self._rk_avail(mode), filename) pipe.sadd(self._rk_tier(mode, tier), filename) if mode == "geminicli" and preview: pipe.sadd(self._rk_preview(mode), filename) else: pipe.srem(self._rk_preview(mode), filename) await pipe.execute() except Exception as e: log.warning(f"Redis sync_cred error: {e}") async def _get_next_available_from_redis( self, mode: str, model_name: Optional[str], exclude_free_tier: bool = False, preview_only: bool = False ) -> Optional[tuple]: """ Redis 快速路径：随机取候选凭证，跳过冷却中的，返回 (filename, credential_data)。 失败或池为空时返回 None，由调用方降级到 MongoDB。 """ try: # 选择候选池优先级：preview_only > exclude_free_tier > 全量池 if preview_only and exclude_free_tier: # preview 且非 free：preview ∩ (pro ∪ ultra) preview_set = await self._redis.smembers(self._rk_preview(mode)) pro_members = await self._redis.smembers(self._rk_tier(mode, "pro")) ultra_members = await self._redis.smembers(self._rk_tier(mode, "ultra")) non_free = pro_members | ultra_members all_candidates = list(preview_set & non_free) if not all_candidates: log.debug(f"[Redis MISS] mode={mode} preview+non-free: no candidates, fallback to MongoDB") return None sample_size = min(len(all_candidates), 10) candidates = random.sample(all_candidates, sample_size) elif preview_only: preview_key = self._rk_preview(mode) preview_size = await self._redis.scard(preview_key) if preview_size == 0: log.debug(f"[Redis MISS] mode={mode} preview_only: pool empty, fallback to MongoDB") return None sample_size = min(preview_size, 10) candidates = await self._redis.srandmember(preview_key, sample_size) if not candidates: return None elif exclude_free_tier: pro_members = await self._redis.smembers(self._rk_tier(mode, "pro")) ultra_members = await self._redis.smembers(self._rk_tier(mode, "ultra")) all_candidates = list(pro_members | ultra_members) if not all_candidates: log.debug(f"[Redis MISS] mode={mode} exclude_free: no non-free creds, fallback to MongoDB") return None sample_size = min(len(all_candidates), 10) candidates = random.sample(all_candidates, sample_size) else: pool_key = self._rk_avail(mode) pool_size = await self._redis.scard(pool_key) if pool_size == 0: log.debug(f"[Redis MISS] mode={mode} pool_key={pool_key}: pool empty, fallback to MongoDB") return None sample_size = min(pool_size, 10) candidates = await self._redis.srandmember(pool_key, sample_size) if not candidates: return None # 过滤冷却中的凭证 if model_name: escaped = self._escape_model_name(model_name) for filename in candidates: cd_key = self._rk_cd(mode, filename, escaped) if not await self._redis.exists(cd_key): credential_data = await self.get_credential(filename, mode) if mode == "antigravity": state = await self.get_credential_state(filename, mode) credential_data = credential_data or {} credential_data["enable_credit"] = bool(state.get("enable_credit", False)) log.debug(f"[Redis HIT] mode={mode} model={model_name} -> {filename}") return filename, credential_data # 所有候选都在冷却中，降级到 MongoDB log.debug(f"[Redis MISS] mode={mode} model={model_name}: all {len(candidates)} candidates in cooldown, fallback to MongoDB") return None else: filename = candidates[0] credential_data = await self.get_credential(filename, mode) if mode == "antigravity": state = await self.get_credential_state(filename, mode) credential_data = credential_data or {} credential_data["enable_credit"] = bool(state.get("enable_credit", False)) log.debug(f"[Redis HIT] mode={mode} -> {filename}") return filename, credential_data except Exception as e: log.warning(f"Redis get_next_available error: {e}") return None async def close(self) -> None: """关闭 MongoDB 连接""" if self._redis: await self._redis.aclose() self._redis = None self._redis_enabled = False if self._client: self._client.close() self._client = None self._db = None self._initialized = False log.debug("MongoDB storage closed") def _ensure_initialized(self): """确保已初始化""" if not self._initialized: raise RuntimeError("MongoDB manager not initialized") def _get_collection_name(self, mode: str) -> str: """根据 mode 获取对应的集合名""" if mode == "antigravity": return "antigravity_credentials" elif mode == "geminicli": return "credentials" else: raise ValueError(f"Invalid mode: {mode}. Must be 'geminicli' or 'antigravity'") # ============ SQL 方法 ============ async def get_next_available_credential( self, mode: str = "geminicli", model_name: Optional[str] = None ) -> Optional[tuple[str, Dict[str, Any]]]: """ 随机获取一个可用凭证（负载均衡） - 未禁用 - 如果提供了 model_name，还会检查模型级冷却 - 随机选择 Args: mode: 凭证模式 ("geminicli" 或 "antigravity") model_name: 完整模型名（如 "gemini-2.0-flash-exp"） Note: - 开启 Redis 时：利用 Redis Set 随机选凭证 + TTL key 判断冷却 - 未开启 Redis 时：使用 count + random skip + limit(1) """ self._ensure_initialized() # Redis 快速路径：根据模型名派生过滤标志，直接在 Redis 分桶中筛选 if self._redis_enabled: model_lower = model_name.lower() if model_name else "" exclude_free = mode == "geminicli" and "pro" in model_lower preview_only = mode == "geminicli" and "preview" in model_lower result = await self._get_next_available_from_redis( mode, model_name, exclude_free_tier=exclude_free, preview_only=preview_only ) if result is not None: return result # result 为 None：池为空或所有候选都冷却中，降级到 MongoDB 以扩大样本空间 log.debug(f"[MongoDB fallback] mode={mode} model={model_name}") try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] current_time = time.time() # 构建普通查询（避免 $sample 聚合导致全集合扫描） match_query: Dict[str, Any] = {"disabled": False} # pro 模型只允许非 free tier 凭证 if mode == "geminicli" and model_name and "pro" in model_name.lower(): match_query["tier"] = {"$ne": "free"} # preview 模型只允许 preview=True 的凭证 if mode == "geminicli" and model_name and "preview" in model_name.lower(): match_query["preview"] = True # 冷却检查：直接用 MongoDB 查询表达，无需 $addFields if model_name: escaped_model_name = self._escape_model_name(model_name) field = f"model_cooldowns.{escaped_model_name}" match_query["$or"] = [ {field: {"$exists": False}}, {field: {"$lte": current_time}}, ] # 统计符合条件的凭证总数（走索引，极快） count = await collection.count_documents(match_query) if count == 0: return None # 随机偏移 + limit(1)，替代 $sample，避免全集合随机排序 skip_n = random.randint(0, count - 1) projection = {"filename": 1, "credential_data": 1, "enable_credit": 1, "_id": 0} docs = await collection.find(match_query, projection).skip(skip_n).limit(1).to_list(1) if docs: doc = docs[0] credential_data = doc.get("credential_data") or {} if mode == "antigravity": credential_data["enable_credit"] = bool(doc.get("enable_credit", False)) return doc["filename"], credential_data return None except Exception as e: log.error(f"Error getting next available credential (mode={mode}, model_name={model_name}): {e}") return None async def get_available_credentials_list(self, mode: str = "geminicli") -> List[str]: """ 获取所有可用凭证列表 - 未禁用 - 按轮换顺序排序 """ self._ensure_initialized() try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] pipeline = [ {"$match": {"disabled": False}}, {"$sort": {"rotation_order": 1}}, {"$project": {"filename": 1, "_id": 0}} ] docs = await collection.aggregate(pipeline).to_list(length=None) return [doc["filename"] for doc in docs] except Exception as e: log.error(f"Error getting available credentials list (mode={mode}): {e}") return [] # ============ StorageBackend 协议方法 ============ async def store_credential(self, filename: str, credential_data: Dict[str, Any], mode: str = "geminicli") -> bool: """存储或更新凭证""" self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] current_ts = time.time() # 使用 upsert + $setOnInsert # 如果文档存在，只更新 credential_data 和 updated_at # 如果文档不存在，设置所有默认字段 # 先尝试更新现有文档 result = await collection.update_one( {"filename": filename}, { "$set": { "credential_data": credential_data, "updated_at": current_ts, } } ) # 如果没有匹配到（新凭证），需要插入 if result.matched_count == 0: # 获取下一个 rotation_order pipeline = [ {"$group": {"_id": None, "max_order": {"$max": "$rotation_order"}}}, {"$project": {"_id": 0, "next_order": {"$add": ["$max_order", 1]}}} ] result_list = await collection.aggregate(pipeline).to_list(length=1) next_order = result_list[0]["next_order"] if result_list else 0 # 插入新凭证（使用 insert_one，因为我们已经确认不存在） try: new_credential = { "filename": filename, "credential_data": credential_data, "disabled": False, "error_codes": [], "error_messages": [], "last_success": current_ts, "user_email": None, "model_cooldowns": {}, "preview": True, "tier": "pro", "rotation_order": next_order, "call_count": 0, "created_at": current_ts, "updated_at": current_ts, } if mode == "antigravity": new_credential["enable_credit"] = False await collection.insert_one(new_credential) # 新凭证插入成功，添加到 Redis 可用池 await self._redis_add_cred(mode, filename) except Exception as insert_error: # 处理并发插入导致的重复键错误 if "duplicate key" in str(insert_error).lower(): # 重试更新（已存在的凭证，无需更新 Redis） await collection.update_one( {"filename": filename}, {"$set": {"credential_data": credential_data, "updated_at": current_ts}} ) else: raise log.debug(f"Stored credential: {filename} (mode={mode})") return True except Exception as e: log.error(f"Error storing credential {filename}: {e}") return False async def get_credential(self, filename: str, mode: str = "geminicli") -> Optional[Dict[str, Any]]: """获取凭证数据""" self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 精确匹配，只投影需要的字段 doc = await collection.find_one( {"filename": filename}, {"credential_data": 1, "_id": 0} ) if doc: return doc.get("credential_data") return None except Exception as e: log.error(f"Error getting credential {filename}: {e}") return None async def list_credentials(self, mode: str = "geminicli") -> List[str]: """列出所有凭证文件名""" self._ensure_initialized() try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 使用聚合管道 pipeline = [ {"$sort": {"rotation_order": 1}}, {"$project": {"filename": 1, "_id": 0}} ] docs = await collection.aggregate(pipeline).to_list(length=None) return [doc["filename"] for doc in docs] except Exception as e: log.error(f"Error listing credentials: {e}") return [] async def delete_credential(self, filename: str, mode: str = "geminicli") -> bool: """删除凭证""" self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 精确匹配删除 result = await collection.delete_one({"filename": filename}) deleted_count = result.deleted_count if deleted_count > 0: # 从 Redis 池中移除 await self._redis_remove_cred(mode, filename) log.debug(f"Deleted {deleted_count} credential(s): {filename} (mode={mode})") return True else: log.warning(f"No credential found to delete: {filename} (mode={mode})") return False except Exception as e: log.error(f"Error deleting credential {filename}: {e}") return False async def get_duplicate_credentials_by_email(self, mode: str = "geminicli") -> Dict[str, Any]: """ 获取按邮箱分组的重复凭证信息（只查询邮箱和文件名，不加载完整凭证数据） 用于去重操作 Args: mode: 凭证模式 ("geminicli" 或 "antigravity") Returns: 包含 email_groups（邮箱分组）、duplicate_count（重复数量）、no_email_count（无邮箱数量）的字典 """ self._ensure_initialized() try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 使用聚合管道，只查询 filename 和 user_email 字段 pipeline = [ { "$project": { "filename": 1, "user_email": 1, "_id": 0 } }, { "$sort": {"filename": 1} } ] docs = await collection.aggregate(pipeline).to_list(length=None) # 按邮箱分组 email_to_files = {} no_email_files = [] for doc in docs: filename = doc.get("filename") user_email = doc.get("user_email") if user_email: if user_email not in email_to_files: email_to_files[user_email] = [] email_to_files[user_email].append(filename) else: no_email_files.append(filename) # 找出重复的邮箱组 duplicate_groups = [] total_duplicate_count = 0 for email, files in email_to_files.items(): if len(files) > 1: # 保留第一个文件，其他为重复 duplicate_groups.append({ "email": email, "kept_file": files[0], "duplicate_files": files[1:], "duplicate_count": len(files) - 1, }) total_duplicate_count += len(files) - 1 return { "email_groups": email_to_files, "duplicate_groups": duplicate_groups, "duplicate_count": total_duplicate_count, "no_email_files": no_email_files, "no_email_count": len(no_email_files), "unique_email_count": len(email_to_files), "total_count": len(docs), } except Exception as e: log.error(f"Error getting duplicate credentials by email: {e}") return { "email_groups": {}, "duplicate_groups": [], "duplicate_count": 0, "no_email_files": [], "no_email_count": 0, "unique_email_count": 0, "total_count": 0, } async def update_credential_state( self, filename: str, state_updates: Dict[str, Any], mode: str = "geminicli" ) -> bool: """更新凭证状态""" self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 过滤只更新状态字段 valid_updates = { k: v for k, v in state_updates.items() if k in self.STATE_FIELDS } if mode != "antigravity": valid_updates.pop("enable_credit", None) if not valid_updates: return True valid_updates["updated_at"] = time.time() # 精确匹配更新 result = await collection.update_one( {"filename": filename}, {"$set": valid_updates} ) updated_count = result.modified_count + result.matched_count # 如果 disabled 发生变化，同步 Redis 池成员关系 if self._redis_enabled and "disabled" in valid_updates: if valid_updates["disabled"]: # 直接禁用：从集合中移除 await self._redis_remove_cred(mode, filename) else: # 重新启用：需要读取当前 tier/preview 以正确放入分桶 doc = await collection.find_one( {"filename": filename}, projection={"tier": 1, "preview": 1, "_id": 0}, ) tier_val = (doc or {}).get("tier", "pro") or "pro" preview_val = (doc or {}).get("preview", True) await self._redis_sync_cred(mode, filename, disabled=False, tier=tier_val, preview=preview_val) elif self._redis_enabled and ("tier" in valid_updates or "preview" in valid_updates): # tier 或 preview 更新：重新同步分桶（只在凭证未禁用时） doc = await collection.find_one( {"filename": filename}, projection={"disabled": 1, "tier": 1, "preview": 1, "_id": 0}, ) if doc and not doc.get("disabled", False): tier_val = doc.get("tier", "pro") or "pro" preview_val = doc.get("preview", True) await self._redis_sync_cred(mode, filename, disabled=False, tier=tier_val, preview=preview_val) return updated_count > 0 except Exception as e: log.error(f"Error updating credential state {filename}: {e}") return False async def get_credential_state(self, filename: str, mode: str = "geminicli") -> Dict[str, Any]: """获取凭证状态（不包含error_messages）""" self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] current_time = time.time() # 精确匹配 doc = await collection.find_one({"filename": filename}) if doc: model_cooldowns = doc.get("model_cooldowns", {}) # 过滤掉损坏的数据(dict类型)和过期的冷却 if model_cooldowns: model_cooldowns = { k: v for k, v in model_cooldowns.items() if isinstance(v, (int, float)) and v > current_time } state = { "disabled": doc.get("disabled", False), "error_codes": doc.get("error_codes", []), "last_success": doc.get("last_success", current_time), "user_email": doc.get("user_email"), "model_cooldowns": model_cooldowns, "preview": doc.get("preview", True), "tier": doc.get("tier", "pro"), } if mode == "antigravity": state["enable_credit"] = doc.get("enable_credit", False) return state # 返回默认状态 default_state = { "disabled": False, "error_codes": [], "last_success": current_time, "user_email": None, "model_cooldowns": {}, "preview": True, "tier": "pro", } if mode == "antigravity": default_state["enable_credit"] = False return default_state except Exception as e: log.error(f"Error getting credential state {filename}: {e}") return {} async def get_all_credential_states(self, mode: str = "geminicli") -> Dict[str, Dict[str, Any]]: """获取所有凭证状态（不包含error_messages）""" self._ensure_initialized() try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 使用投影只获取需要的字段（不包含error_messages） projection = { "filename": 1, "disabled": 1, "error_codes": 1, "last_success": 1, "user_email": 1, "model_cooldowns": 1, "preview": 1, "tier": 1, "enable_credit": 1, "_id": 0 } cursor = collection.find({}, projection=projection) states = {} current_time = time.time() async for doc in cursor: filename = doc["filename"] model_cooldowns = doc.get("model_cooldowns", {}) # 自动过滤掉已过期的模型CD if model_cooldowns: model_cooldowns = { k: v for k, v in model_cooldowns.items() if isinstance(v, (int, float)) and v > current_time } state = { "disabled": doc.get("disabled", False), "error_codes": doc.get("error_codes", []), "last_success": doc.get("last_success", time.time()), "user_email": doc.get("user_email"), "model_cooldowns": model_cooldowns, "preview": doc.get("preview", True), "tier": doc.get("tier", "pro"), } if mode == "antigravity": state["enable_credit"] = doc.get("enable_credit", False) states[filename] = state return states except Exception as e: log.error(f"Error getting all credential states: {e}") return {} async def get_credentials_summary( self, offset: int = 0, limit: Optional[int] = None, status_filter: str = "all", mode: str = "geminicli", error_code_filter: Optional[str] = None, cooldown_filter: Optional[str] = None, preview_filter: Optional[str] = None, tier_filter: Optional[str] = None ) -> Dict[str, Any]: """ 获取凭证的摘要信息（不包含完整凭证数据）- 支持分页和状态筛选 Args: offset: 跳过的记录数（默认0） limit: 返回的最大记录数（None表示返回所有） status_filter: 状态筛选（all=全部, enabled=仅启用, disabled=仅禁用） mode: 凭证模式 ("geminicli" 或 "antigravity") error_code_filter: 错误码筛选（格式如"400"或"403"，筛选包含该错误码的凭证） cooldown_filter: 冷却状态筛选（"in_cooldown"=冷却中, "no_cooldown"=未冷却） preview_filter: Preview筛选（"preview"=支持preview, "no_preview"=不支持preview，仅geminicli模式有效） tier_filter: tier筛选（"free", "pro", "ultra"） Returns: 包含 items（凭证列表）、total（总数）、offset、limit 的字典 """ self._ensure_initialized() try: # 根据 mode 选择集合名 collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 构建查询条件 query = {} if status_filter == "enabled": query["disabled"] = False elif status_filter == "disabled": query["disabled"] = True # 错误码筛选 - 兼容存储为数字或字符串的情况 if error_code_filter and str(error_code_filter).strip().lower() != "all": filter_value = str(error_code_filter).strip() query_values = [filter_value] try: query_values.append(int(filter_value)) except ValueError: pass query["error_codes"] = {"$in": query_values} # 计算全局统计数据（不受筛选条件影响） global_stats = {"total": 0, "normal": 0, "disabled": 0} stats_pipeline = [ { "$group": { "_id": "$disabled", "count": {"$sum": 1} } } ] stats_result = await collection.aggregate(stats_pipeline).to_list(length=10) for item in stats_result: count = item["count"] global_stats["total"] += count if item["_id"]: global_stats["disabled"] = count else: global_stats["normal"] = count # 获取所有匹配的文档（用于冷却筛选，因为需要在Python中判断） projection = { "filename": 1, "disabled": 1, "error_codes": 1, "last_success": 1, "user_email": 1, "rotation_order": 1, "model_cooldowns": 1, "preview": 1, "tier": 1, "enable_credit": 1, "_id": 0 } cursor = collection.find(query, projection=projection).sort("rotation_order", 1) all_summaries = [] current_time = time.time() async for doc in cursor: model_cooldowns = doc.get("model_cooldowns", {}) # 自动过滤掉已过期的模型CD active_cooldowns = {} if model_cooldowns: active_cooldowns = { k: v for k, v in model_cooldowns.items() if isinstance(v, (int, float)) and v > current_time } summary = { "filename": doc["filename"], "disabled": doc.get("disabled", False), "error_codes": doc.get("error_codes", []), "last_success": doc.get("last_success", current_time), "user_email": doc.get("user_email"), "rotation_order": doc.get("rotation_order", 0), "model_cooldowns": active_cooldowns, "preview": doc.get("preview", True), "tier": doc.get("tier", "pro"), } if mode == "antigravity": summary["enable_credit"] = bool(doc.get("enable_credit", False)) if mode == "geminicli" and preview_filter: preview_value = summary.get("preview", True) if preview_filter == "preview" and not preview_value: continue if preview_filter == "no_preview" and preview_value: continue # 应用tier筛选 if tier_filter and tier_filter in ("free", "pro", "ultra"): if summary["tier"] != tier_filter: continue # 应用冷却筛选 if cooldown_filter == "in_cooldown": # 只保留有冷却的凭证 if active_cooldowns: all_summaries.append(summary) elif cooldown_filter == "no_cooldown": # 只保留没有冷却的凭证 if not active_cooldowns: all_summaries.append(summary) else: # 不筛选冷却状态 all_summaries.append(summary) # 应用分页 total_count = len(all_summaries) if limit is not None: summaries = all_summaries[offset:offset + limit] else: summaries = all_summaries[offset:] return { "items": summaries, "total": total_count, "offset": offset, "limit": limit, "stats": global_stats, } except Exception as e: log.error(f"Error getting credentials summary: {e}") return { "items": [], "total": 0, "offset": offset, "limit": limit, "stats": {"total": 0, "normal": 0, "disabled": 0}, } # ============ 配置管理（内存缓存 + 可选 Redis）============ def _rk_config(self, key: str) -> str: """配置项的 Redis key""" return f"gcli:config:{key}" def _rk_config_all(self) -> str: """所有配置的 Redis Hash key""" return "gcli:config" async def _load_config_to_redis(self) -> None: """将所有配置从 MongoDB 同步到 Redis Hash""" if not self._redis_enabled: return try: config_collection = self._db["config"] cursor = config_collection.find({}) mapping = {} async for doc in cursor: mapping[doc["key"]] = json.dumps(doc.get("value")) pipe = self._redis.pipeline() pipe.delete(self._rk_config_all()) if mapping: pipe.hset(self._rk_config_all(), mapping=mapping) await pipe.execute() log.debug(f"Synced {len(mapping)} config items to Redis") except Exception as e: log.warning(f"Failed to sync config to Redis: {e}") async def set_config(self, key: str, value: Any) -> bool: """设置配置（写入数据库；Redis 启用时写 Redis，否则更新内存缓存）""" self._ensure_initialized() try: config_collection = self._db["config"] await config_collection.update_one( {"key": key}, {"$set": {"value": value, "updated_at": time.time()}}, upsert=True, ) if self._redis_enabled: try: await self._redis.hset(self._rk_config_all(), key, json.dumps(value)) except Exception as e: log.warning(f"Redis config set error for key={key}: {e}") else: self._config_cache[key] = value return True except Exception as e: log.error(f"Error setting config {key}: {e}") return False async def reload_config_cache(self): """重新加载配置缓存（在批量修改配置后调用）""" self._ensure_initialized() if self._redis_enabled: await self._load_config_to_redis() else: self._config_loaded = False await self._load_config_cache() log.info("Config cache reloaded from database") async def get_config(self, key: str, default: Any = None) -> Any: """获取配置（Redis 启用时从 Redis 读取，否则从内存缓存）""" self._ensure_initialized() if self._redis_enabled: try: raw = await self._redis.hget(self._rk_config_all(), key) if raw is not None: return json.loads(raw) return default except Exception as e: log.warning(f"Redis config get error for key={key}: {e}") return default return self._config_cache.get(key, default) async def get_all_config(self) -> Dict[str, Any]: """获取所有配置（Redis 启用时从 Redis 读取，否则从内存缓存）""" self._ensure_initialized() if self._redis_enabled: try: raw_map = await self._redis.hgetall(self._rk_config_all()) return {k: json.loads(v) for k, v in raw_map.items()} except Exception as e: log.warning(f"Redis config getall error: {e}") return {} return self._config_cache.copy() async def delete_config(self, key: str) -> bool: """删除配置""" self._ensure_initialized() try: config_collection = self._db["config"] result = await config_collection.delete_one({"key": key}) if self._redis_enabled: try: await self._redis.hdel(self._rk_config_all(), key) except Exception as e: log.warning(f"Redis config delete error for key={key}: {e}") else: self._config_cache.pop(key, None) return result.deleted_count > 0 except Exception as e: log.error(f"Error deleting config {key}: {e}") return False async def get_credential_errors(self, filename: str, mode: str = "geminicli") -> Dict[str, Any]: """ 专门获取凭证的错误信息（包含 error_codes 和 error_messages） Args: filename: 凭证文件名 mode: 凭证模式 ("geminicli" 或 "antigravity") Returns: 包含 error_codes 和 error_messages 的字典 """ self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 精确匹配 doc = await collection.find_one( {"filename": filename}, {"error_codes": 1, "error_messages": 1, "_id": 0} ) if doc: return { "filename": filename, "error_codes": doc.get("error_codes", []), "error_messages": doc.get("error_messages", []), } # 凭证不存在，返回空错误信息 return { "filename": filename, "error_codes": [], "error_messages": [], } except Exception as e: log.error(f"Error getting credential errors {filename}: {e}") return { "filename": filename, "error_codes": [], "error_messages": [], "error": str(e) } # ============ 模型级冷却管理 ============ async def set_model_cooldown( self, filename: str, model_name: str, cooldown_until: Optional[float], mode: str = "geminicli" ) -> bool: """ 设置特定模型的冷却时间 Args: filename: 凭证文件名 model_name: 模型名（完整模型名，如 "gemini-2.0-flash-exp"） cooldown_until: 冷却截止时间戳（None 表示清除冷却） mode: 凭证模式 ("geminicli" 或 "antigravity") Returns: 是否成功 """ self._ensure_initialized() # 统一使用 basename 处理文件名 filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] # 转义模型名中的点号 escaped_model_name = self._escape_model_name(model_name) # 使用原子操作直接更新，避免竞态条件 if cooldown_until is None: # 删除指定模型的冷却 result = await collection.update_one( {"filename": filename}, { "$unset": {f"model_cooldowns.{escaped_model_name}": ""}, "$set": {"updated_at": time.time()} } ) else: # 设置冷却时间 result = await collection.update_one( {"filename": filename}, { "$set": { f"model_cooldowns.{escaped_model_name}": cooldown_until, "updated_at": time.time() } } ) if result.matched_count == 0: log.warning(f"Credential {filename} not found") return False # 同步写入 Redis TTL key if self._redis_enabled: cd_key = self._rk_cd(mode, filename, escaped_model_name) if cooldown_until is None: await self._redis.delete(cd_key) else: ttl = int(cooldown_until - time.time()) if ttl > 0: await self._redis.setex(cd_key, ttl, str(cooldown_until)) else: # 冷却已经过期，确保清除 await self._redis.delete(cd_key) log.debug(f"Set model cooldown: {filename}, model_name={model_name}, cooldown_until={cooldown_until}") return True except Exception as e: log.error(f"Error setting model cooldown for {filename}: {e}") return False async def clear_all_model_cooldowns( self, filename: str, mode: str = "geminicli" ) -> bool: """ 清除某个凭证的所有模型冷却时间 Args: filename: 凭证文件名 mode: 凭证模式 ("geminicli" 或 "antigravity") Returns: 是否成功 """ self._ensure_initialized() filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] doc = await collection.find_one( {"filename": filename}, {"model_cooldowns": 1, "_id": 0} ) if not doc: log.warning(f"Credential {filename} not found") return False model_cooldowns = doc.get("model_cooldowns") or {} await collection.update_one( {"filename": filename}, { "$set": { "model_cooldowns": {}, "updated_at": time.time(), } } ) if self._redis_enabled and isinstance(model_cooldowns, dict) and model_cooldowns: redis_keys = [self._rk_cd(mode, filename, escaped_model) for escaped_model in model_cooldowns.keys()] await self._redis.delete(*redis_keys) log.debug(f"Cleared all model cooldowns: {filename} (mode={mode})") return True except Exception as e: log.error(f"Error clearing all model cooldowns for {filename}: {e}") return False async def record_success( self, filename: str, model_name: Optional[str] = None, mode: str = "geminicli" ) -> None: """ 成功调用后的条件写入： - 只有当前 error_codes 非空时才清除错误并写 last_success - 只有当前存在该模型的冷却键时才清除 通过 MongoDB 服务端条件匹配实现 """ self._ensure_initialized() filename = os.path.basename(filename) try: collection_name = self._get_collection_name(mode) collection = self._db[collection_name] now = time.time() # 条件写入：只有 error_codes 非空时才触发，避免无意义的写 IO await collection.update_one( {"filename": filename, "error_codes": {"$ne": []}}, {"$set": { "last_success": now, "error_codes": [], "error_messages": {}, "updated_at": now, }} ) # 条件删除模型冷却：只有该键存在时才写入 if model_name: escaped = self._escape_model_name(model_name) await collection.update_one( {"filename": filename, f"model_cooldowns.{escaped}": {"$exists": True}}, {"$unset": {f"model_cooldowns.{escaped}": ""}, "$set": {"updated_at": now}} ) # 同步删除 Redis 冷却 key if self._redis_enabled: await self._redis.delete(self._rk_cd(mode, filename, escaped)) except Exception as e: log.error(f"Error recording success for {filename}: {e}")