ForcePilot/backend/package/yuxi/channels/adapters/qqbot/reconnect.py

from __future__ import annotations

import asyncio
import logging
import random
import time
from enum import Enum, auto
from collections.abc import Callable, Awaitable

from yuxi.channels.adapters.qqbot.constants import ECode, SERVER_CLOSE_CODE_MAP

logger = logging.getLogger(__name__)

_RAPID_DISCONNECT_THRESHOLD_S = 5.0
_RAPID_DISCONNECT_MAX_WARNINGS = 3


class ReconnectState(Enum):
    DISCONNECTED = auto()
    CONNECTING = auto()
    CONNECTED = auto()
    RECONNECTING = auto()
    IDENTIFYING = auto()
    RESUMING = auto()
    BACKOFF = auto()
    FROZEN = auto()


class CloseCodeCategory(Enum):
    ABNORMAL = auto()
    RECOVERABLE = auto()
    FATAL = auto()
    SERVER_SIDE = auto()
    SERVER_ERROR = auto()
    RATE_LIMITED = auto()


class ServerErrorCategory(Enum):
    OVERLOAD = auto()
    MAINTENANCE = auto()
    NETWORK = auto()
    INTERNAL = auto()
    UNAVAILABLE = auto()
    TIMEOUT = auto()
    UNKNOWN = auto()


_RECOVERABLE_CODES: set[int] = {
    ECode.UNKNOWN_ERROR,
    ECode.UNKNOWN_OPCODE,
    ECode.DECODE_ERROR,
    ECode.NOT_AUTHENTICATED,
    ECode.AUTHENTICATION_FAILED,
    ECode.RATE_LIMITED,
    4009,
    ECode.INVALID_INTENT,
    ECode.INVALID_SHARD,
}

_FATAL_CODES: set[int] = {
    ECode.INVALID_API_VERSION,
    ECode.INVALID_SEQ,
    4013,
    ECode.INVALID_SHARD_COUNT,
    ECode.BOT_REMOVED,
    ECode.ACCOUNT_BANNED,
}

_ABNORMAL_CODES: set[int] = {4000, 4008, 4011}

_RATELIMIT_CODES: set[int] = {ECode.RATE_LIMITED, 4009}

_SERVER_ERROR_CODES: set[int] = set(range(4900, 4914))

_SERVER_OVERLOAD_CODES: set[int] = {4901, 4904, 4912}
_SERVER_MAINTENANCE_CODES: set[int] = {4902, 4905}
_SERVER_NETWORK_CODES: set[int] = {4903, 4906, 4907}
_SERVER_UNAVAILABLE_CODES: set[int] = {4908, 4909}
_SERVER_TIMEOUT_CODES: set[int] = {4913}
_SERVER_INTERNAL_CODES: set[int] = {4900, 4910, 4911}


def classify_server_error_category(code: int) -> ServerErrorCategory:
    if code in _SERVER_OVERLOAD_CODES:
        return ServerErrorCategory.OVERLOAD
    if code in _SERVER_MAINTENANCE_CODES:
        return ServerErrorCategory.MAINTENANCE
    if code in _SERVER_NETWORK_CODES:
        return ServerErrorCategory.NETWORK
    if code in _SERVER_UNAVAILABLE_CODES:
        return ServerErrorCategory.UNAVAILABLE
    if code in _SERVER_TIMEOUT_CODES:
        return ServerErrorCategory.TIMEOUT
    if code in _SERVER_INTERNAL_CODES:
        return ServerErrorCategory.INTERNAL
    return ServerErrorCategory.UNKNOWN


def get_server_error_name(code: int) -> str:
    return SERVER_CLOSE_CODE_MAP.get(code, f"server_error_{code}")


def classify_close_code(code: int | None) -> CloseCodeCategory:
    if code is None:
        return CloseCodeCategory.ABNORMAL
    if code in _RATELIMIT_CODES:
        return CloseCodeCategory.RATE_LIMITED
    if code in _FATAL_CODES:
        return CloseCodeCategory.FATAL
    if code in _RECOVERABLE_CODES:
        return CloseCodeCategory.RECOVERABLE
    if code in _ABNORMAL_CODES:
        return CloseCodeCategory.ABNORMAL
    if code in _SERVER_ERROR_CODES:
        return CloseCodeCategory.SERVER_ERROR
    return CloseCodeCategory.SERVER_SIDE if 4000 <= code < 5000 else CloseCodeCategory.ABNORMAL


class QQBotReconnectManager:
    def __init__(
        self,
        base_delay: float = 1.0,
        max_delay: float = 60.0,
        jitter: float = 0.3,
        max_retries: int = 10,
        resume_timeout: float = 15.0,
    ) -> None:
        self._state = ReconnectState.DISCONNECTED
        self._base_delay = base_delay
        self._max_delay = max_delay
        self._jitter = jitter
        self._max_retries = max_retries
        self._resume_timeout = resume_timeout
        self._retry_count = 0
        self._session_id: str | None = None
        self._last_seq: int | None = None
        self._seq_reset_lock = asyncio.Lock()
        self._state_listeners: list[Callable[[ReconnectState, ReconnectState], Awaitable[None]]] = []
        self._last_connect_time: float = 0.0
        self._rapid_disconnect_count: int = 0
        self._last_disconnect_code: int | None = None
        self._last_disconnect_time: float = 0.0

    @property
    def state(self) -> ReconnectState:
        return self._state

    @property
    def session_id(self) -> str | None:
        return self._session_id

    @property
    def last_seq(self) -> int | None:
        return self._last_seq

    def add_state_listener(self, listener: Callable[[ReconnectState, ReconnectState], Awaitable[None]]) -> None:
        self._state_listeners.append(listener)

    async def _notify_state_change(self, old: ReconnectState, new: ReconnectState) -> None:
        for listener in self._state_listeners:
            try:
                await listener(old, new)
            except Exception:
                logger.exception("ReconnectManager state listener error")

    async def transition(self, new: ReconnectState) -> None:
        old = self._state
        if old == new:
            return
        self._state = new
        logger.info("ReconnectManager: %s -> %s", old.name, new.name)
        await self._notify_state_change(old, new)

    def on_identify_success(self, session_id: str) -> None:
        self._session_id = session_id
        self._retry_count = 0
        self._last_connect_time = time.monotonic()

    def mark_connected(self) -> None:
        self._last_connect_time = time.monotonic()
        self._rapid_disconnect_count = 0

    async def record_seq(self, seq: int) -> None:
        async with self._seq_reset_lock:
            self._last_seq = seq

    def seq_reset(self) -> None:
        self._last_seq = None

    def on_hello(self) -> None:
        pass

    def should_resume(self) -> bool:
        return self._session_id is not None and self._last_seq is not None

    async def on_disconnect(self, code: int | None) -> None:
        now = time.monotonic()
        self._last_disconnect_code = code
        self._last_disconnect_time = now

        category = classify_close_code(code)
        self._check_rapid_disconnect(code, category, now)

        if category == CloseCodeCategory.FATAL:
            logger.error("ReconnectManager: fatal close code %s, freezing", code)
            await self.transition(ReconnectState.FROZEN)
            return

        if category == CloseCodeCategory.SERVER_ERROR:
            error_name = get_server_error_name(code) if code else "unknown"
            sub_category = classify_server_error_category(code) if code else ServerErrorCategory.UNKNOWN
            logger.warning(
                "ReconnectManager: server error code=%s name=%s category=%s",
                code,
                error_name,
                sub_category.name,
            )

            self._retry_count += 1
            if self._retry_count >= self._max_retries:
                logger.error("ReconnectManager: server error retries exhausted, freezing")
                await self.transition(ReconnectState.FROZEN)
                return

            delay = self._calc_delay_for_server_error(sub_category)
            logger.warning(
                "ReconnectManager: server error close code %s (%s), backing off %.1fs (retry %d/%d)",
                code,
                error_name,
                delay,
                self._retry_count,
                self._max_retries,
            )
            self.seq_reset()
            self._session_id = None
            await self.transition(ReconnectState.BACKOFF)
            await asyncio.sleep(delay)
            await self.transition(ReconnectState.IDENTIFYING)
            return

        if category == CloseCodeCategory.RATE_LIMITED:
            self._retry_count += 1
            if self._retry_count > 3:
                logger.error("ReconnectManager: rate-limited retries exhausted, freezing")
                await self.transition(ReconnectState.FROZEN)
                return
            self.seq_reset()

        if self._retry_count >= self._max_retries:
            logger.error("ReconnectManager: max retries (%d) exhausted", self._max_retries)
            await self.transition(ReconnectState.FROZEN)
            return

        self._retry_count += 1
        delay = self._calc_delay()
        logger.info("ReconnectManager: backing off %.1fs (retry %d/%d)", delay, self._retry_count, self._max_retries)
        await self.transition(ReconnectState.BACKOFF)
        await asyncio.sleep(delay)

        if self.should_resume():
            await self.transition(ReconnectState.RESUMING)
        else:
            self.seq_reset()
            await self.transition(ReconnectState.IDENTIFYING)

    def _calc_delay(self) -> float:
        raw = min(self._base_delay * (2 ** (self._retry_count - 1)), self._max_delay)
        jittered = raw * (1 + random.uniform(-self._jitter, self._jitter))
        return max(0.5, min(jittered, self._max_delay))

    def _calc_delay_for_server_error(self, sub_category: ServerErrorCategory) -> float:
        base = self._calc_delay()
        if sub_category == ServerErrorCategory.OVERLOAD:
            return base * 3.0
        if sub_category == ServerErrorCategory.MAINTENANCE:
            return base * 5.0
        if sub_category == ServerErrorCategory.NETWORK:
            return base * 1.5
        if sub_category == ServerErrorCategory.TIMEOUT:
            return base * 1.5
        return base * 2.0

    def _check_rapid_disconnect(self, code: int | None, category: CloseCodeCategory, now: float) -> None:
        if category == CloseCodeCategory.FATAL:
            return
        if self._last_connect_time == 0:
            return
        elapsed = now - self._last_connect_time
        if elapsed > _RAPID_DISCONNECT_THRESHOLD_S:
            self._rapid_disconnect_count = 0
            return
        self._rapid_disconnect_count += 1
        if self._rapid_disconnect_count >= _RAPID_DISCONNECT_MAX_WARNINGS:
            logger.error(
                "ReconnectManager: RAPID DISCONNECT LOOP DETECTED - "
                "%d disconnects within %.1fs threshold, last code=%s category=%s elapsed=%.2fs",
                self._rapid_disconnect_count,
                _RAPID_DISCONNECT_THRESHOLD_S,
                code,
                category.name,
                elapsed,
            )
        else:
            logger.warning(
                "ReconnectManager: rapid disconnect #%d/%d within %.1fs threshold, code=%s category=%s elapsed=%.2fs",
                self._rapid_disconnect_count,
                _RAPID_DISCONNECT_MAX_WARNINGS,
                _RAPID_DISCONNECT_THRESHOLD_S,
                code,
                category.name,
                elapsed,
            )

    async def reset(self) -> None:
        self._retry_count = 0
        self._session_id = None
        self._last_seq = None
        self._last_connect_time = 0.0
        self._rapid_disconnect_count = 0
        self._last_disconnect_code = None
        self._last_disconnect_time = 0.0
        await self.transition(ReconnectState.DISCONNECTED)