diff --git a/backend/package/yuxi/channels/channel_shared/__init__.py b/backend/package/yuxi/channels/channel_shared/__init__.py new file mode 100644 index 00000000..91081604 --- /dev/null +++ b/backend/package/yuxi/channels/channel_shared/__init__.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from .multi_account_base import ( + AccountConfigBase, + AccountLifecycleMixin, + ConnectionPoolMixin, + MultiAccountBase, + MultiAccountRegistryBase, +) + +__all__ = [ + "AccountConfigBase", + "AccountLifecycleMixin", + "ConnectionPoolMixin", + "MultiAccountBase", + "MultiAccountRegistryBase", +] diff --git a/backend/package/yuxi/channels/channel_shared/multi_account_base.py b/backend/package/yuxi/channels/channel_shared/multi_account_base.py new file mode 100644 index 00000000..039d3602 --- /dev/null +++ b/backend/package/yuxi/channels/channel_shared/multi_account_base.py @@ -0,0 +1,146 @@ +from __future__ import annotations + +import asyncio +import logging +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + + +@dataclass +class AccountConfigBase: + account_id: str + label: str = "" + enabled: bool = True + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def configured(self) -> bool: + return True + + +class MultiAccountRegistryBase[T](ABC): + accounts: dict[str, T] + default_account_id: str + + @abstractmethod + def resolve_account(self, account_id: str | None = None) -> T | None: ... + + @abstractmethod + def list_account_ids(self) -> list[str]: ... + + @abstractmethod + def list_enabled(self) -> list[T]: ... + + @property + def is_multi(self) -> bool: + return len(self.list_enabled()) > 1 + + @property + def account_count(self) -> int: + return len(self.accounts) + + +class AccountLifecycleMixin: + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if not hasattr(cls, "_lifecycle_hooks"): + cls._lifecycle_hooks: dict[str, list] = { + "on_account_added": [], + "on_account_removed": [], + "on_account_enabled": [], + "on_account_disabled": [], + "on_account_config_changed": [], + } + + def register_hook(self, event: str, handler) -> None: + hooks = getattr(self, "_lifecycle_hooks", {}) + if event in hooks: + hooks[event].append(handler) + + async def _fire_hook(self, event: str, account_id: str, data: dict[str, Any]) -> None: + hooks = getattr(self, "_lifecycle_hooks", {}) + for handler in hooks.get(event, []): + try: + result = handler(account_id, data) + if asyncio.iscoroutine(result): + await result + except Exception: + logger.debug("Lifecycle hook error for event '%s'", event, exc_info=True) + + async def add_account(self, account_id: str, config: dict[str, Any]) -> bool: + raise NotImplementedError + + async def remove_account(self, account_id: str) -> bool: + raise NotImplementedError + + async def delete_account(self, account_id: str) -> bool: + return await self.remove_account(account_id) + + async def update_account(self, account_id: str, config: dict[str, Any]) -> bool: + raise NotImplementedError + + +class ConnectionPoolMixin: + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if not hasattr(cls, "_connections"): + cls._connections: dict[str, Any] = {} + if not hasattr(cls, "_connect_tasks"): + cls._connect_tasks: dict[str, asyncio.Task] = {} + if not hasattr(cls, "_ready_events"): + cls._ready_events: dict[str, asyncio.Event] = {} + + def get_connection(self, account_id: str) -> Any | None: + return self._connections.get(account_id) + + def is_connected(self, account_id: str) -> bool: + ready_evt = self._ready_events.get(account_id) + return ready_evt.is_set() if ready_evt else False + + @property + def active_ids(self) -> list[str]: + return [aid for aid, evt in self._ready_events.items() if evt.is_set()] + + async def disconnect_account(self, account_id: str) -> bool: + conn = self._connections.pop(account_id, None) + task = self._connect_tasks.pop(account_id, None) + ready_evt = self._ready_events.pop(account_id, None) + disconnected = False + + if task and not task.done(): + task.cancel() + try: + await task + except (asyncio.CancelledError, Exception): + pass + disconnected = True + + if conn is not None: + try: + close_fn = getattr(conn, "close", None) or getattr(conn, "disconnect", None) + if close_fn: + result = close_fn() + if asyncio.iscoroutine(result): + await result + except Exception: + pass + disconnected = True + + if ready_evt: + ready_evt.clear() + + return disconnected + + async def disconnect_all(self) -> None: + for account_id in list(self._connections.keys()): + await self.disconnect_account(account_id) + self._connections.clear() + self._connect_tasks.clear() + self._ready_events.clear() + + +class MultiAccountBase(AccountLifecycleMixin, ConnectionPoolMixin): + pass