ForcePilot/backend/package/yuxi/channel/extensions/yuanbao/client.py

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from typing import Any

from yuxi.channel.extensions.yuanbao.codec.biz_codec import BizCodec
from yuxi.channel.extensions.yuanbao.codec.conn_codec import ConnCodec
from yuxi.channel.extensions.yuanbao.types import GatewayState

logger = logging.getLogger(__name__)

try:
    import websockets
    from websockets.asyncio.client import ClientConnection

    HAS_WEBSOCKETS = True
except ImportError:
    HAS_WEBSOCKETS = False
    logger.warning("websockets package not installed, Yuanbao WS disabled")


class YuanbaoWsClient:
    HEARTBEAT_INTERVAL = 30.0
    HEARTBEAT_TIMEOUT = 10.0

    def __init__(
        self,
        ws_url: str,
        access_token: str,
        on_message: Callable[[bytes], Any] | None = None,
    ):
        if not HAS_WEBSOCKETS:
            raise ImportError(
                "websockets package is required for Yuanbao WebSocket. "
                "Install it with: pip install websockets>=13.0"
            )

        self._ws_url = ws_url
        self._access_token = access_token
        self._on_message = on_message
        self._connection: ClientConnection | None = None
        self._state = GatewayState.DISCONNECTED
        self._heartbeat_task: asyncio.Task | None = None
        self._listen_task: asyncio.Task | None = None
        self._cancel_event: asyncio.Event | None = None
        self._receive_lock = asyncio.Lock()

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

    @property
    def is_connected(self) -> bool:
        return self._state == GatewayState.READY and self._connection is not None

    async def connect(self, cancel_event: asyncio.Event) -> None:
        if not HAS_WEBSOCKETS:
            raise ImportError("websockets package required")

        self._cancel_event = cancel_event
        self._state = GatewayState.CONNECTING

        url = f"{self._ws_url}?access_token={self._access_token}"

        self._connection = await websockets.connect(
            url,
            ping_interval=None,
            max_size=2 ** 23,
            close_timeout=5.0,
        )

        auth_frame = ConnCodec.encode_auth_req(self._access_token)
        await self._connection.send(auth_frame)

        auth_raw = await asyncio.wait_for(self._connection.recv(), timeout=15.0)
        auth_response = ConnCodec.decode_frame(auth_raw)

        if not ConnCodec.is_auth_success(auth_response):
            self._state = GatewayState.AUTH_FAILED
            await self._connection.close()
            self._connection = None
            raise ConnectionError(f"Auth failed: {auth_response}")

        self._state = GatewayState.READY

        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        self._listen_task = asyncio.create_task(self._listen_loop())

    async def disconnect(self) -> None:
        self._state = GatewayState.DISCONNECTED

        if self._heartbeat_task:
            self._heartbeat_task.cancel()
            self._heartbeat_task = None
        if self._listen_task:
            self._listen_task.cancel()
            self._listen_task = None

        if self._connection:
            try:
                await self._connection.close()
            except Exception:
                pass
            self._connection = None

    async def send_binary(self, data: bytes) -> None:
        if not self._connection or self._state != GatewayState.READY:
            raise ConnectionError("Not connected")
        async with self._receive_lock:
            await self._connection.send(data)

    async def send_text_msg(self, msg: dict[str, Any]) -> None:
        encoded = BizCodec.encode_message(msg)
        await self.send_binary(encoded)

    async def _heartbeat_loop(self) -> None:
        while self._state == GatewayState.READY and not (
            self._cancel_event and self._cancel_event.is_set()
        ):
            try:
                await asyncio.sleep(self.HEARTBEAT_INTERVAL)
                if self._connection and self._state == GatewayState.READY:
                    ping_frame = ConnCodec.encode_ping()
                    await self._connection.send(ping_frame)
                    pong_raw = await asyncio.wait_for(
                        self._connection.recv(), timeout=self.HEARTBEAT_TIMEOUT
                    )
                    pong_frame = ConnCodec.decode_frame(pong_raw)
                    if not ConnCodec.is_ping(pong_frame):
                        logger.warning("Unexpected heartbeat response: %s", pong_frame)
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("Heartbeat error, disconnecting")
                self._state = GatewayState.ERROR
                break

    async def _listen_loop(self) -> None:
        while self._state == GatewayState.READY and self._connection and not (
            self._cancel_event and self._cancel_event.is_set()
        ):
            try:
                raw = await self._connection.recv()
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("WS receive error")
                self._state = GatewayState.ERROR
                break

            try:
                frame = ConnCodec.decode_frame(raw)
                if ConnCodec.is_ping(frame):
                    pong = ConnCodec.encode_pong()
                    await self._connection.send(pong)
                elif ConnCodec.is_data_frame(frame):
                    if self._on_message:
                        data_payload = ConnCodec.extract_data_payload(frame)
                        if data_payload:
                            await self._on_message(data_payload)
                else:
                    logger.debug("Unhandled frame type: %s", frame.get("type"))
            except Exception:
                logger.exception("Error processing inbound WS frame")