from __future__ import annotations import asyncio from dataclasses import dataclass, field from typing import Any from slack_sdk.socket_mode.aiohttp import SocketModeClient from slack_sdk.web.async_client import AsyncWebClient from yuxi.utils.logging_config import logger @dataclass class SlackAccountInfo: account_id: str team_id: str = "" team_name: str = "" bot_token: str = "" app_token: str = "" signing_secret: str = "" bot_user_id: str = "" scopes: list[str] = field(default_factory=list) metadata: dict[str, Any] = field(default_factory=dict) @property def has_socket_mode(self) -> bool: return bool(self.app_token) @property def has_http_mode(self) -> bool: return bool(self.signing_secret) @dataclass class ConnectedAccount: account: SlackAccountInfo client: AsyncWebClient socket_handler: SocketModeClient | None = None socket_task: asyncio.Task | None = None connected: bool = False async def disconnect(self) -> None: self.connected = False if self.socket_handler: try: self.socket_handler.disconnect() except Exception: pass if self.socket_task and not self.socket_task.done(): self.socket_task.cancel() try: await self.socket_task except (asyncio.CancelledError, Exception): pass self.socket_handler = None self.socket_task = None @dataclass class MultiAccountConnections: connections: dict[str, ConnectedAccount] = field(default_factory=dict) def get_connection(self, account_id: str) -> ConnectedAccount | None: return self.connections.get(account_id) def get_client(self, account_id: str) -> AsyncWebClient | None: conn = self.connections.get(account_id) return conn.client if conn else None def is_connected(self, account_id: str) -> bool: conn = self.connections.get(account_id) return conn.connected if conn else False @property def active_ids(self) -> list[str]: return [aid for aid, c in self.connections.items() if c.connected] async def disconnect_all(self) -> None: for conn in list(self.connections.values()): await conn.disconnect() self.connections.clear() async def connect_account(account: SlackAccountInfo) -> ConnectedAccount: client = AsyncWebClient(token=account.bot_token) connected = ConnectedAccount(account=account, client=client) try: auth = await client.auth_test() if not auth.get("ok"): logger.error(f"auth.test failed for account {account.account_id}: {auth.get('error')}") return connected connected.account.bot_user_id = auth.get("user_id", "") connected.account.team_id = auth.get("team_id", "") if account.app_token: socket_handler = SocketModeClient( app_token=account.app_token, web_client=client, auto_reconnect_enabled=True, ) connected.socket_handler = socket_handler connected.socket_task = asyncio.create_task(socket_handler.connect_async()) connected.connected = True logger.info(f"Slack account {account.account_id} connected via Socket Mode") elif account.signing_secret: connected.connected = True logger.info(f"Slack account {account.account_id} connected via HTTP Mode") else: connected.connected = True logger.info(f"Slack account {account.account_id} connected (client only)") return connected except Exception as e: logger.error(f"Failed to connect account {account.account_id}: {e}") return connected @dataclass class MultiAccountRegistry: accounts: dict[str, SlackAccountInfo] = field(default_factory=dict) default_account_id: str = "" top_level_config: dict[str, Any] = field(default_factory=dict) def get_account(self, account_id: str | None = None) -> SlackAccountInfo | None: if not account_id: account_id = self.default_account_id return self.accounts.get(account_id) def get_merged_config(self, account_id: str | None = None) -> dict[str, Any]: account = self.get_account(account_id) merged = dict(self.top_level_config) if account: account_config = account.metadata.get("config", {}) or {} merged.update(account_config) return merged def register_account(self, account: SlackAccountInfo) -> None: self.accounts[account.account_id] = account if not self.default_account_id: self.default_account_id = account.account_id @classmethod def from_config(cls, config: dict[str, Any] | None) -> MultiAccountRegistry: if not config: return cls() accounts_config = config.get("accounts", []) or [] top_level_config = dict(config) top_level_config.pop("accounts", None) if isinstance(accounts_config, dict): accounts = {} default_id = "" for acc_id, acc_config in accounts_config.items(): if default_id == "": default_id = acc_id if isinstance(acc_config, dict): accounts[acc_id] = SlackAccountInfo( account_id=acc_id, team_id=acc_config.get("team_id", ""), team_name=acc_config.get("team_name", ""), bot_token=acc_config.get("bot_token", ""), app_token=acc_config.get("app_token", ""), signing_secret=acc_config.get("signing_secret", ""), bot_user_id=acc_config.get("bot_user_id", ""), scopes=list(acc_config.get("scopes", [])), metadata=dict(acc_config.get("metadata", {})), ) return cls(accounts=accounts, default_account_id=default_id, top_level_config=top_level_config) accounts = {} default_id = "" for acc_config in accounts_config: if not isinstance(acc_config, dict): continue acc_id = acc_config.get("id", "") if default_id == "": default_id = acc_id accounts[acc_id] = SlackAccountInfo( account_id=acc_id, team_id=acc_config.get("team_id", ""), team_name=acc_config.get("team_name", ""), bot_token=acc_config.get("bot_token", ""), app_token=acc_config.get("app_token", ""), signing_secret=acc_config.get("signing_secret", ""), bot_user_id=acc_config.get("bot_user_id", ""), scopes=list(acc_config.get("scopes", [])), metadata=dict(acc_config.get("metadata", {})), ) return cls(accounts=accounts, default_account_id=default_id, top_level_config=top_level_config) async def _init_account_connection( account: SlackAccountInfo, ) -> dict[str, Any] | None: try: from slack_sdk.web.async_client import AsyncWebClient client = AsyncWebClient(token=account.bot_token) auth = await client.auth_test() if auth.get("ok"): return { "connected": True, "team_id": auth.get("team_id", account.team_id), "bot_user_id": auth.get("user_id", account.bot_user_id), } return {"connected": False, "error": auth.get("error")} except Exception as e: logger.error(f"Failed to init account {account.account_id}: {e}") return {"connected": False, "error": str(e)}