from __future__ import annotations import asyncio import time from collections import defaultdict, deque from typing import Any from sqlalchemy import func, select from yuxi.channels.base import BaseChannelAdapter from yuxi.channels.infra.broadcast import EventBroadcaster from yuxi.channels.infra.circuit_breaker import CircuitBreaker from yuxi.channels.infra.config_watcher import ConfigWatcher from yuxi.channels.services.context import GatewayRequestContext from yuxi.channels.services.doctor import ConfigDoctor, DiagnosisIssue from yuxi.channels.exceptions import ChannelException from yuxi.channels.services.maintenance import MaintenanceRunner from yuxi.channels.models import ChannelStatus from yuxi.channels.registry import ChannelRegistry, _BUILTIN_ADAPTERS, _load_builtin_adapters from yuxi.channels.router import MessageRouter from yuxi.channels.services.runtime_state import RuntimeState from yuxi.channels.services.stats_collector import StatsCollector from yuxi.channels.services.webhook_registry import WebhookRegistry from yuxi.channels.services.ws_logger import WsLogger from yuxi.storage.postgres.manager import pg_manager from yuxi.storage.postgres.models_channels import ChannelMsgRecord from yuxi.utils.logging_config import logger HEALTH_CHECK_INTERVAL = 60 CONFIG_WATCH_INTERVAL = 30.0 class ChannelManager: def __init__(self, registry: ChannelRegistry | None = None, router: MessageRouter | None = None): self._registry = registry or ChannelRegistry() self._router = router or MessageRouter(channel_manager=self) self._adapters: dict[str, BaseChannelAdapter] = {} self._circuit_breakers: dict[str, CircuitBreaker] = defaultdict(lambda: CircuitBreaker()) self._health_tasks: dict[str, asyncio.Task] = {} self._initialized = False self._phase: str = "not_started" self._rate_limiters: dict[str, deque[float]] = {} self._rate_limit_locks: dict[str, asyncio.Lock] = {} self._channels_config: dict[str, dict[str, Any]] = {} self._watcher: ConfigWatcher | None = None self._doctor: ConfigDoctor | None = None self._broadcaster: EventBroadcaster | None = None self._ctx: GatewayRequestContext | None = None self.runtime_state = RuntimeState() self._auth_limiter: Any = None self._ws_handlers: dict[str, Any] = {} self._scheduled_tasks: list[asyncio.Task] = [] self._ws_logger: WsLogger | None = None self._maintenance_runner: MaintenanceRunner | None = None self._stats_collector: StatsCollector | None = None self._webhook_registry: WebhookRegistry | None = None @property def phase(self) -> str: return self._phase @property def context(self) -> GatewayRequestContext | None: return self._ctx @property def broadcaster(self) -> EventBroadcaster | None: return self._broadcaster @property def watcher(self) -> ConfigWatcher | None: return self._watcher @property def doctor(self) -> ConfigDoctor | None: return self._doctor async def load_config(self) -> None: if self._phase not in ("not_started",): return _load_builtin_adapters() for channel_id, adapter_cls in _BUILTIN_ADAPTERS.items(): self._registry.register(channel_id, adapter_cls) from yuxi import config as conf self._channels_config = getattr(conf, "channels", {}) self._phase = "config_loaded" logger.info("ChannelManager: config loaded") async def prepare_bootstrap(self) -> None: if self._phase not in ("config_loaded",): return async with pg_manager.get_async_session_context() as db: await self._ensure_virtual_department(db) await self._ensure_default_agent_config(db) self._phase = "bootstrapped" logger.info("ChannelManager: bootstrap prepared") async def start_channels(self) -> None: if self._phase not in ("config_loaded", "bootstrapped"): return for channel_id in self._registry.list_channels(): channel_conf = self._channels_config.get(channel_id, {}) if channel_conf.get("enabled", False): try: await self.start_channel(channel_id, channel_conf) except Exception: logger.exception(f"Failed to start channel {channel_id}") self._phase = "channels_started" logger.info(f"ChannelManager: channels started: {list(self._adapters.keys())}") async def start_subscriptions(self) -> None: if self._phase not in ("channels_started",): return self._phase = "fully_running" logger.info("ChannelManager: subscriptions started, fully running") async def initialize(self) -> None: if self._initialized: return _load_builtin_adapters() for channel_id, adapter_cls in _BUILTIN_ADAPTERS.items(): self._registry.register(channel_id, adapter_cls) from yuxi import config as conf self._channels_config = getattr(conf, "channels", {}) for channel_id in self._registry.list_channels(): channel_conf = self._channels_config.get(channel_id, {}) if channel_conf.get("enabled", False): try: await self.start_channel(channel_id, channel_conf) except Exception: logger.exception(f"Failed to start channel {channel_id}") async with pg_manager.get_async_session_context() as db: await self._ensure_virtual_department(db) await self._ensure_default_agent_config(db) self._initialized = True self._phase = "fully_running" logger.info(f"ChannelManager initialized with channels: {list(self._adapters.keys())}") async def shutdown(self) -> None: for channel_id in list(self._adapters.keys()): try: await self.stop_channel(channel_id) except Exception: logger.exception(f"Failed to stop channel {channel_id} during shutdown") for task in self._health_tasks.values(): task.cancel() for task in self._scheduled_tasks: task.cancel() if self._watcher: await self._watcher.stop() self._initialized = False logger.info("ChannelManager shutdown complete") async def startup(self) -> None: if self._initialized: return try: await self._stage_load_config() await self._stage_prepare_bootstrap() await self._stage_start_early_runtime() await self._stage_create_auth_limiter() await self._stage_init_channels() await self._stage_create_runtime_state() await self._stage_start_runtime_services() await self._stage_activate_scheduled_services() await self._stage_attach_ws_handlers() await self._stage_start_event_subscriptions() self._initialized = True self._phase = "fully_running" logger.info("ChannelManager: 10-stage startup complete") except Exception: logger.exception(f"ChannelManager startup failed at phase: {self._phase}") raise async def _stage_load_config(self) -> None: if self._phase not in ("not_started",): return _load_builtin_adapters() for channel_id, adapter_cls in _BUILTIN_ADAPTERS.items(): self._registry.register(channel_id, adapter_cls) from yuxi import config as conf self._channels_config = getattr(conf, "channels", {}) self._phase = "config_loaded" logger.info("ChannelManager: [1/10] config loaded") async def _stage_prepare_bootstrap(self) -> None: if self._phase not in ("config_loaded",): return async with pg_manager.get_async_session_context() as db: await self._ensure_virtual_department(db) await self._ensure_default_agent_config(db) self._phase = "bootstrapped" logger.info("ChannelManager: [2/10] bootstrap prepared") async def _stage_start_early_runtime(self) -> None: if self._phase not in ("bootstrapped",): return self._broadcaster = EventBroadcaster() self._ctx = GatewayRequestContext( runtime_config=self._channels_config, start_channel=self._ctx_start_channel, stop_channel=self._ctx_stop_channel, mark_channel_logged_out=self._ctx_mark_channel_logged_out, broadcast_fn=self._broadcaster.broadcast, node_send_to_session_fn=self._broadcaster.node_send_to_session, ) self._phase = "early_runtime" logger.info("ChannelManager: [3/10] early runtime started") async def _stage_create_auth_limiter(self) -> None: if self._phase not in ("early_runtime",): return self._phase = "auth_limiter_created" logger.info("ChannelManager: [4/10] auth rate limiter created") async def _stage_init_channels(self) -> None: if self._phase not in ("early_runtime", "auth_limiter_created"): return for channel_id in self._registry.list_channels(): channel_conf = self._channels_config.get(channel_id, {}) if channel_conf.get("enabled", False): try: await self.start_channel(channel_id, channel_conf) except Exception: logger.exception(f"Failed to start channel {channel_id}") self._phase = "channels_started" logger.info(f"ChannelManager: [5/10] channels started: {list(self._adapters.keys())}") async def _stage_create_runtime_state(self) -> None: if self._phase not in ("channels_started",): return self.runtime_state = RuntimeState( started_at=time.monotonic(), active_channels=len(self._adapters), phase=self._phase, node_id="forcepilot-gateway", main_node=True, services=["doctor", "maintenance", "stats", "webhooks"], ) self._ws_logger = WsLogger(max_entries=1000) self._phase = "runtime_state_created" logger.info("ChannelManager: [6/10] runtime state created") async def _stage_start_runtime_services(self) -> None: if self._phase not in ("runtime_state_created",): return self._doctor = ConfigDoctor(self) self._maintenance_runner = MaintenanceRunner(self, self.runtime_state) self._stats_collector = StatsCollector(self.runtime_state) self._webhook_registry = WebhookRegistry(self) self._phase = "runtime_services_started" logger.info("ChannelManager: [7/10] runtime services started (doctor+maintenance+stats+webhooks)") async def _stage_activate_scheduled_services(self) -> None: if self._phase not in ("runtime_services_started",): return self._watcher = ConfigWatcher(self) await self._watcher.watch(interval=CONFIG_WATCH_INTERVAL) self._scheduled_tasks.append(asyncio.create_task(self._maintenance_runner.run())) self._scheduled_tasks.append(asyncio.create_task(self._stats_collector.run())) self._scheduled_tasks.append(asyncio.create_task(self._webhook_registry.run())) self._phase = "scheduled_services_active" logger.info("ChannelManager: [8/10] scheduled services activated (watcher+maintenance+stats+webhooks)") async def _stage_attach_ws_handlers(self) -> None: if self._phase not in ("scheduled_services_active",): return if self._broadcaster: self._broadcaster.subscribe_callback("channel.status_change", self._on_channel_status_change) self._broadcaster.subscribe_callback("tick", self._on_tick) self._broadcaster.subscribe_callback("chat", self._on_chat_event) self._broadcaster.subscribe_callback("channel.logout", self._on_channel_logout) self._broadcaster.subscribe_callback("config.reload", self._on_config_reload) self._phase = "ws_handlers_attached" logger.info("ChannelManager: [9/10] websocket handlers attached") async def _stage_start_event_subscriptions(self) -> None: if self._phase not in ("ws_handlers_attached",): return self.runtime_state.services = ["doctor", "maintenance", "stats", "webhooks", "watcher", "ws"] self._phase = "subscriptions_started" logger.info("ChannelManager: [10/10] event subscriptions started") async def _on_channel_status_change(self, event: str, payload: Any) -> None: logger.debug(f"WS event: {event}") async def _on_tick(self, event: str, payload: Any) -> None: pass async def _on_chat_event(self, event: str, payload: Any) -> None: logger.debug(f"WS chat event: {event}") async def _on_channel_logout(self, event: str, payload: Any) -> None: channel_id = payload.get("channel_id") if isinstance(payload, dict) else None account_id = payload.get("account_id") if isinstance(payload, dict) else None logger.info(f"Channel logout: {channel_id}/{account_id}") async def _on_config_reload(self, event: str, payload: Any) -> None: await self.reload_config_now() async def _ctx_start_channel(self, channel_id: str, config: dict) -> None: await self.start_channel(channel_id, config) async def _ctx_stop_channel(self, channel_id: str) -> None: await self.stop_channel(channel_id) async def _ctx_mark_channel_logged_out(self, channel_id: str, account_id: str) -> None: logger.info(f"Channel {channel_id} account {account_id} marked as logged out") async def diagnose_channels(self) -> list[DiagnosisIssue]: if self._doctor is None: return [] return await self._doctor.diagnose() async def auto_fix_channel(self, issue: DiagnosisIssue) -> bool: if self._doctor is None: return False return await self._doctor.auto_fix(issue) async def reload_config_now(self, changed_keys: list[str] | None = None) -> dict[str, Any]: if self._watcher is None: return {} return await self._watcher.reload_now(changed_keys) async def start_channel(self, channel_id: str, config: dict[str, Any] | None = None) -> None: if channel_id in self._adapters: logger.warning(f"Channel {channel_id} already running") return adapter_cls = self._registry.get(channel_id) if not adapter_cls: raise ValueError(f"No adapter registered for channel {channel_id}") config = config or self._channels_config.get(channel_id, {}) adapter = adapter_cls(config=config) adapter.on_message(self._handle_inbound_message) pre_connect_result = await adapter.pre_connect() if pre_connect_result: qr_url = pre_connect_result.get("qr_url") if qr_url: logger.info(f"Channel {channel_id} requires QR scan: {qr_url}") await adapter.connect() self._adapters[channel_id] = adapter if channel_id not in self._health_tasks: self._health_tasks[channel_id] = asyncio.create_task(self._health_check_loop(channel_id)) logger.info(f"Channel {channel_id} started") async def stop_channel(self, channel_id: str) -> None: adapter = self._adapters.pop(channel_id, None) if not adapter: return task = self._health_tasks.pop(channel_id, None) if task: task.cancel() try: await task except asyncio.CancelledError: pass try: await adapter.disconnect() except Exception: logger.exception(f"Error disconnecting channel {channel_id}") logger.info(f"Channel {channel_id} stopped") async def restart_channel(self, channel_id: str) -> None: config = self._channels_config.get(channel_id, {}) await self.stop_channel(channel_id) await self.start_channel(channel_id, config) async def send_outbound(self, channel_id: str, response) -> None: adapter = self._adapters.get(channel_id) if not adapter: raise ChannelException(f"Channel {channel_id} not found", retryable=False) cb = self._circuit_breakers[channel_id] await cb.call(lambda: adapter.send(response)) async def get_channel_status(self, channel_id: str | None = None) -> dict: if channel_id: return await self._get_single_channel_status(channel_id) all_channels = {} for cid in self._registry.list_channels(): info = await self._get_single_channel_status(cid) stats = info.get("stats") or {} all_channels[cid] = { "channel_id": cid, "channel_type": info.get("channel_type"), "display_name": info.get("display_name"), "enabled": info.get("enabled", False), "status": info.get("status", "not_found"), "total_count": stats.get("total_messages", 0), "today_count": stats.get("today_messages", 0), "active_connections": 1 if info.get("status") == ChannelStatus.CONNECTED.value else 0, "health": info.get("health"), } return {"channels": all_channels} async def update_channel_config(self, channel_id: str, config_updates: dict[str, Any]) -> dict: adapter = self._adapters.get(channel_id) if not adapter: raise ChannelException(f"Channel {channel_id} not found", retryable=False) if hasattr(adapter, "reload_config"): await adapter.reload_config(config_updates) else: for key, value in config_updates.items(): adapter.config[key] = value if channel_id not in self._channels_config: self._channels_config[channel_id] = {} self._channels_config[channel_id].update(config_updates) return {"channel_id": channel_id, "config_updated": True} async def test_channel(self, channel_id: str) -> dict: adapter = self._adapters.get(channel_id) if not adapter: return {"channel_id": channel_id, "test_result": "failure", "error": "Channel not running"} try: start = time.monotonic() health = await adapter.health_check() latency_ms = (time.monotonic() - start) * 1000 return { "channel_id": channel_id, "test_result": "success" if health.status == "healthy" else "degraded", "latency_ms": round(latency_ms, 1), "health": health.model_dump(), } except Exception as e: return {"channel_id": channel_id, "test_result": "failure", "error": str(e)} async def check_rate_limit(self, key: str, max_req: int, window_seconds: int) -> bool: lock = self._rate_limit_locks.setdefault(key, asyncio.Lock()) async with lock: now = time.monotonic() history = self._rate_limiters.setdefault(key, deque(maxlen=max_req)) while history and now - history[0] > window_seconds: history.popleft() if len(history) >= max_req: return False history.append(now) return True async def _handle_inbound_message(self, message) -> None: try: await self._router.route_inbound(message) except Exception: logger.exception("Error handling inbound message") async def _health_check_loop(self, channel_id: str) -> None: while channel_id in self._adapters: await asyncio.sleep(HEALTH_CHECK_INTERVAL) adapter = self._adapters.get(channel_id) if not adapter: break cb = self._circuit_breakers.get(channel_id) try: health = await adapter.health_check() if cb and health.status == "healthy": try: await cb.record_success() except Exception: pass elif cb: try: await cb.record_failure() except Exception: pass except asyncio.CancelledError: break except Exception as e: logger.warning(f"Health check failed for {channel_id}: {e}") if cb: try: await cb.record_failure() except Exception: pass async def _get_single_channel_status(self, channel_id: str) -> dict: adapter = self._adapters.get(channel_id) if not adapter: adapter_cls = self._registry.get(channel_id) if adapter_cls: caps = ( adapter_cls.capabilities.model_dump() if hasattr(adapter_cls, "capabilities") else { "text_chunk_limit": adapter_cls.text_chunk_limit, "supports_markdown": adapter_cls.supports_markdown, "supports_streaming": adapter_cls.supports_streaming, "max_media_size_mb": adapter_cls.max_media_size_mb, } ) channel_type = adapter_cls.channel_type.value return { "channel_id": channel_id, "channel_type": channel_type, "display_name": None, "enabled": False, "status": ChannelStatus.DISABLED.value, "config": {"enabled": False}, "capabilities": caps, "health": None, "stats": None, } return {"channel_id": channel_id, "status": "not_found"} health = None try: health_result = await adapter.health_check() health = health_result.model_dump() except Exception: pass caps = ( type(adapter).capabilities.model_dump() if hasattr(type(adapter), "capabilities") else { "text_chunk_limit": type(adapter).text_chunk_limit, "supports_markdown": type(adapter).supports_markdown, "supports_streaming": type(adapter).supports_streaming, "max_media_size_mb": type(adapter).max_media_size_mb, } ) circuit_state = self._circuit_breakers.get(channel_id) cb_state = circuit_state.state.value if circuit_state else "unknown" channel_type = adapter.channel_type.value display_name = adapter.config.get("display_name") return { "channel_id": channel_id, "channel_type": channel_type, "display_name": display_name, "enabled": adapter.config.get("enabled", False) if hasattr(adapter, "config") else True, "status": self._adapter_status(adapter), "config": adapter.config, "capabilities": caps, "health": health, "circuit_state": cb_state, "stats": await self._get_channel_stats(channel_id), } async def _get_channel_stats(self, channel_id: str) -> dict: try: from yuxi.utils.datetime_utils import utc_now_naive now = utc_now_naive() today_start = now.replace(hour=0, minute=0, second=0, microsecond=0) async with pg_manager.get_async_session_context() as session: result = await session.execute( select( func.count().label("total"), func.sum( func.cast( (ChannelMsgRecord.status == "success").cast(func.Integer), func.Integer, ) ).label("success_count"), func.sum( func.cast( (ChannelMsgRecord.status == "error").cast(func.Integer), func.Integer, ) ).label("error_count"), ).where(ChannelMsgRecord.channel_id == channel_id) ) row = result.one() total = row.total or 0 success = row.success_count or 0 error = row.error_count or 0 today_result = await session.execute( select(func.count()).where( ChannelMsgRecord.channel_id == channel_id, ChannelMsgRecord.created_at >= today_start, ) ) today = today_result.scalar() or 0 return { "total_messages": total, "today_messages": today, "success_count": int(success), "error_count": int(error), "success_rate": round(success / total, 3) if total > 0 else 0, } except Exception: return { "total_messages": 0, "today_messages": 0, "success_count": 0, "error_count": 0, "success_rate": 0, } def _adapter_status(self, adapter: BaseChannelAdapter) -> str: return getattr(adapter, "status", "unknown") async def _ensure_virtual_department(self, db) -> None: from yuxi.storage.postgres.models_business import Department result = await db.execute(select(Department).where(Department.id == -1)) if result.scalar_one_or_none() is None: dept = Department( id=-1, name="\u6e20\u9053\u7f51\u5173", description="\u591a\u6e20\u9053\u7f51\u5173\u865a\u62df\u90e8\u95e8", ) db.add(dept) await db.commit() logger.info("Created virtual department (id=-1)") async def _ensure_default_agent_config(self, db) -> None: from yuxi.repositories.agent_config_repository import AgentConfigRepository repo = AgentConfigRepository(db) await repo.get_or_create_default( department_id=-1, agent_id="ChatbotAgent", created_by="system", ) logger.info("Ensured default agent config for ChatbotAgent") channel_manager = ChannelManager()