from __future__ import annotations import asyncio import logging import random import time from collections.abc import Callable from dataclasses import dataclass, field logger = logging.getLogger(__name__) @dataclass class AccountConfig: account_id: str app_id: str app_secret: str label: str = "" weight: int = 1 priority: int = 0 group_ids: list[str] = field(default_factory=list) user_ids: list[str] = field(default_factory=list) cooldown_s: float = 30.0 _fail_count: int = 0 _last_fail: float = 0.0 _last_used: float = 0.0 @property def is_cooling_down(self) -> bool: if self._last_fail <= 0: return False return time.time() - self._last_fail < self.cooldown_s def record_success(self) -> None: self._fail_count = 0 self._last_fail = 0.0 self._last_used = time.time() def record_failure(self) -> None: self._fail_count += 1 self._last_fail = time.time() def matches_chat(self, group_id: str = "", user_id: str = "") -> bool: if self.group_ids or self.user_ids: if group_id and self.group_ids and group_id not in self.group_ids: return False if user_id and self.user_ids and user_id not in self.user_ids: return False return True @dataclass class AccountRouteResult: account: AccountConfig account_id: str resolved: bool = True reason: str = "" class MultiAccountManager: def __init__( self, accounts: list[AccountConfig] | None = None, default_rotation_strategy: str = "weighted_round_robin", ): self._accounts: dict[str, AccountConfig] = {} self._rotation_index = 0 self._lock = asyncio.Lock() self._strategy = default_rotation_strategy self._route_fn: Callable | None = None if accounts: for acc in accounts: self._accounts[acc.account_id] = acc @classmethod def from_config(cls, config: dict | None) -> MultiAccountManager: if not config: return cls() accounts_cfg = config.get("accounts", []) if not accounts_cfg: app_id = config.get("app_id", "") app_secret = config.get("app_secret", "") if app_id and app_secret: acc = AccountConfig( account_id="default", app_id=app_id, app_secret=app_secret, label="Default", ) return cls(accounts=[acc]) return cls() accounts = [] for ac in accounts_cfg: accounts.append(AccountConfig( account_id=ac.get("account_id", str(random.randint(1000, 9999))), app_id=ac.get("app_id", ""), app_secret=ac.get("app_secret", ""), label=ac.get("label", ""), weight=ac.get("weight", 1), priority=ac.get("priority", 0), group_ids=ac.get("group_ids", []), user_ids=ac.get("user_ids", []), cooldown_s=ac.get("cooldown_s", 30.0), )) return cls(accounts=accounts) @property def account_count(self) -> int: return len(self._accounts) def get_account(self, account_id: str) -> AccountConfig | None: return self._accounts.get(account_id) async def route( self, group_id: str = "", user_id: str = "", strategy: str | None = None, ) -> AccountRouteResult: async with self._lock: strategy = strategy or self._strategy if self._route_fn is not None: result = self._route_fn(self._accounts, group_id, user_id) if result: return result candidates = [ acc for acc in self._accounts.values() if acc.matches_chat(group_id, user_id) and not acc.is_cooling_down ] if not candidates: all_accounts = [ acc for acc in self._accounts.values() if acc.matches_chat(group_id, user_id) ] if all_accounts: acc = all_accounts[0] return AccountRouteResult( account=acc, account_id=acc.account_id, reason="all cooling down, picked first", ) return AccountRouteResult( account=AccountConfig(account_id="", app_id="", app_secret=""), account_id="", resolved=False, reason="no matching accounts", ) if strategy == "weighted_random": weights = [acc.weight for acc in candidates] total = sum(weights) if total <= 0: acc = candidates[0] else: r = random.uniform(0, total) agg = 0 acc = candidates[0] for candidate in candidates: agg += candidate.weight if r <= agg: acc = candidate break elif strategy == "least_used": acc = min(candidates, key=lambda a: a._last_used) elif strategy == "priority": candidates.sort(key=lambda a: (-a.priority, a._fail_count)) acc = candidates[0] else: idx = self._rotation_index % len(candidates) acc = candidates[idx] self._rotation_index += 1 acc.record_success() return AccountRouteResult( account=acc, account_id=acc.account_id, ) def set_route_fn(self, fn: Callable | None) -> None: self._route_fn = fn async def mark_failure(self, account_id: str) -> None: async with self._lock: acc = self._accounts.get(account_id) if acc: acc.record_failure() logger.warning( "MultiAccount: account %s failed (count=%d)", account_id, acc._fail_count, ) async def mark_success(self, account_id: str) -> None: async with self._lock: acc = self._accounts.get(account_id) if acc: acc.record_success() async def all_cooling_down(self) -> bool: async with self._lock: return all(acc.is_cooling_down for acc in self._accounts.values()) if self._accounts else False def list_accounts(self) -> list[dict]: return [ { "account_id": acc.account_id, "label": acc.label, "weight": acc.weight, "priority": acc.priority, "is_cooling_down": acc.is_cooling_down, "fail_count": acc._fail_count, } for acc in self._accounts.values() ]