from __future__ import annotations import asyncio import json import logging import time from typing import Any, Callable import websockets import websockets.exceptions from yuxi.channel.extensions.qqbot.api_client import QQBotApiClient from yuxi.channel.extensions.qqbot.errors import classify_close_code from yuxi.channel.extensions.qqbot.reconnect import ReconnectStateMachine from yuxi.channel.extensions.qqbot.session import SessionStore from yuxi.channel.extensions.qqbot.types import ( FULL_INTENTS, GatewayEvent, QQBotOpCode, SessionState, ) logger = logging.getLogger(__name__) class QQBotGatewayConnection: RECONNECT_OP = 7 def __init__( self, api_client: QQBotApiClient, session_store: SessionStore, account_id: str = "default", dispatch_handler: Callable[[GatewayEvent], Any] | None = None, reconnect_handler: Callable[[], Any] | None = None, ): self._api_client = api_client self._session_store = session_store self._account_id = account_id self._dispatch_handler = dispatch_handler self._reconnect_handler = reconnect_handler self._ws: websockets.WebSocketClientProtocol | None = None self._session_id: str | None = None self._last_seq: int | None = None self._heartbeat_interval: int = 30000 self._heartbeat_task: asyncio.Task | None = None self._message_task: asyncio.Task | None = None self._connected = False self._cancel_event = asyncio.Event() self._reconnect_state = ReconnectStateMachine() self._bot_openid: str | None = None self._resume_boundary_seq: int | None = None @property def connected(self) -> bool: return self._connected @property def session_id(self) -> str | None: return self._session_id @property def last_seq(self) -> int | None: return self._last_seq @property def bot_openid(self) -> str | None: return self._bot_openid async def start(self) -> None: self._cancel_event.clear() session = self._session_store.load(self._account_id) if session and session.session_id: self._session_id = session.session_id self._last_seq = session.last_seq await self._connect() async def stop(self) -> None: self._cancel_event.set() await self._disconnect() self._save_session() async def _connect(self) -> None: gateway_url = await self._api_client.get_gateway_url() if not gateway_url: raise RuntimeError("Failed to obtain gateway URL") logger.info("[qqbot:%s] Connecting to gateway: %s", self._account_id, gateway_url) try: self._ws = await websockets.connect( gateway_url, ping_interval=None, close_timeout=5, max_size=2**20, ) except Exception as e: raise RuntimeError(f"WebSocket connection failed: {e}") from e self._connected = True self._message_task = asyncio.create_task(self._message_loop(), name=f"qqbot-ws-{self._account_id}") async def _disconnect(self) -> None: self._connected = False await self._cancel_heartbeat() if self._message_task and not self._message_task.done(): self._message_task.cancel() try: await self._message_task except asyncio.CancelledError: pass self._message_task = None if self._ws: try: await self._ws.close(1000) except Exception: pass self._ws = None async def _message_loop(self) -> None: assert self._ws is not None try: async for raw_message in self._ws: if self._cancel_event.is_set(): break try: data = json.loads(raw_message) event = GatewayEvent.from_dict(data) await self._handle_event(event) except json.JSONDecodeError: logger.warning("[qqbot:%s] Invalid JSON message", self._account_id) except websockets.exceptions.ConnectionClosed as e: self._connected = False should_retry, delay = classify_close_code(e.code) if should_retry: logger.warning("[qqbot:%s] Connection closed: code=%d, reason=%s", self._account_id, e.code, e.reason) if self._reconnect_handler: await self._reconnect_handler() else: logger.error("[qqbot:%s] Fatal connection close: code=%d", self._account_id, e.code) raise except Exception: self._connected = False logger.exception("[qqbot:%s] Message loop error", self._account_id) if self._reconnect_handler and not self._cancel_event.is_set(): await self._reconnect_handler() async def _handle_event(self, event: GatewayEvent) -> None: if event.s is not None: self._last_seq = event.s if event.op == QQBotOpCode.HELLO: await self._on_hello(event) elif event.op == QQBotOpCode.HEARTBEAT_ACK: logger.debug("[qqbot:%s] Heartbeat ACK", self._account_id) elif event.op == QQBotOpCode.RECONNECT: logger.info("[qqbot:%s] Server requested reconnect", self._account_id) if self._reconnect_handler: await self._reconnect_handler() elif event.op == QQBotOpCode.INVALID_SESSION: logger.warning("[qqbot:%s] Invalid session, clearing session", self._account_id) self._session_id = None self._last_seq = None if self._reconnect_handler: await self._reconnect_handler() elif event.op == QQBotOpCode.DISPATCH: await self._on_dispatch(event) async def _on_hello(self, event: GatewayEvent) -> None: d = event.d or {} self._heartbeat_interval = d.get("heartbeat_interval", 30000) logger.info( "[qqbot:%s] HELLO received, heartbeat_interval=%dms", self._account_id, self._heartbeat_interval, ) await self._send_identify_or_resume() self._start_heartbeat() async def _send_identify_or_resume(self) -> None: assert self._ws is not None try: token_str = await self._api_client._token_manager.get_token() except Exception as e: logger.error("[qqbot:%s] Failed to get token for identify: %s", self._account_id, e) raise token = f"QQBot {token_str}" if self._session_id is not None and self._last_seq is not None: self._resume_boundary_seq = self._last_seq payload = { "op": QQBotOpCode.RESUME, "d": { "token": token, "session_id": self._session_id, "seq": self._last_seq, }, } logger.info("[qqbot:%s] Sending RESUME: session_id=%s, boundary_seq=%d", self._account_id, self._session_id, self._resume_boundary_seq) else: payload = { "op": QQBotOpCode.IDENTIFY, "d": { "token": token, "intents": FULL_INTENTS, "shard": [0, 1], }, } logger.info("[qqbot:%s] Sending IDENTIFY", self._account_id) await self._ws.send(json.dumps(payload, separators=(",", ":"))) def _start_heartbeat(self) -> None: if self._heartbeat_task and not self._heartbeat_task.done(): return interval = max(self._heartbeat_interval, 1000) / 1000.0 self._heartbeat_task = asyncio.create_task(self._heartbeat_loop(interval), name=f"qqbot-hb-{self._account_id}") async def _cancel_heartbeat(self) -> None: if self._heartbeat_task and not self._heartbeat_task.done(): self._heartbeat_task.cancel() try: await self._heartbeat_task except asyncio.CancelledError: pass self._heartbeat_task = None async def _heartbeat_loop(self, interval: float) -> None: while not self._cancel_event.is_set(): try: await asyncio.wait_for(self._cancel_event.wait(), timeout=interval) return except TimeoutError: pass if self._ws and self._connected: try: payload = { "op": QQBotOpCode.HEARTBEAT, "d": self._last_seq, } await self._ws.send(json.dumps(payload, separators=(",", ":"))) logger.debug("[qqbot:%s] Heartbeat sent, seq=%s", self._account_id, self._last_seq) except Exception: logger.exception("[qqbot:%s] Heartbeat send failed", self._account_id) async def _on_dispatch(self, event: GatewayEvent) -> None: event_type = event.t or "" d = event.d or {} if self._resume_boundary_seq is not None and event.s is not None and event.s <= self._resume_boundary_seq: logger.debug( "[qqbot:%s] Skipping replayed event: type=%s, seq=%d <= boundary=%d", self._account_id, event_type, event.s, self._resume_boundary_seq, ) return if event_type == "READY": self._session_id = d.get("session_id") self._bot_openid = d.get("user", {}).get("id") self._resume_boundary_seq = None logger.info( "[qqbot:%s] READY: session_id=%s, bot_openid=%s", self._account_id, self._session_id, self._bot_openid, ) self._save_session() elif event_type == "RESUMED": logger.info("[qqbot:%s] RESUMED", self._account_id) self._resume_boundary_seq = None if self._dispatch_handler: try: await self._dispatch_handler(event) except Exception: logger.exception("[qqbot:%s] Dispatch handler error for event %s", self._account_id, event_type) def _save_session(self) -> None: state = SessionState( session_id=self._session_id, last_seq=self._last_seq, last_connected_at=time.time(), account_id=self._account_id, saved_at=time.time(), ) self._session_store.save(self._account_id, state)