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

from __future__ import annotations

import asyncio
import logging

from yuxi.channel.extensions.yuanbao.accounts import ResolvedYuanbaoAccount, TokenCache
from yuxi.channel.extensions.yuanbao.client import YuanbaoWsClient
from yuxi.channel.extensions.yuanbao.codec.biz_codec import BizCodec
from yuxi.channel.extensions.yuanbao.inbound.dispatcher import InboundDispatcher
from yuxi.channel.extensions.yuanbao.types import YuanbaoAccountConfig

logger = logging.getLogger(__name__)


class YuanbaoGateway:
    MAX_RECONNECT_ATTEMPTS = 100

    def __init__(
        self,
        account: ResolvedYuanbaoAccount,
        token_cache: TokenCache,
        queue: asyncio.Queue | None = None,
        cancel_event: asyncio.Event | None = None,
        logger: logging.Logger | None = None,
    ):
        self._account = account
        self._token_cache = token_cache
        self._queue = queue
        self._cancel_event = cancel_event or asyncio.Event()
        self._log = logger or logging.getLogger(__name__)
        self._ws_client: YuanbaoWsClient | None = None
        self._dispatcher = InboundDispatcher(queue=self._queue)
        self._config: YuanbaoAccountConfig | None = None
        self._reconnect_count = 0
        self._listen_task: asyncio.Task | None = None

    @property
    def is_connected(self) -> bool:
        return self._ws_client is not None and self._ws_client.is_connected

    def configure(self, config: YuanbaoAccountConfig) -> None:
        self._config = config

    async def start(self) -> None:
        self._reconnect_count = 0
        await self._connect()
        self._listen_task = asyncio.create_task(self._reconnect_loop())

    async def _connect(self) -> None:
        access_token = await self._token_cache.get(self._account)

        self._ws_client = YuanbaoWsClient(
            ws_url=self._account.ws_url,
            access_token=access_token,
            on_message=self._on_ws_message,
        )

        try:
            await self._ws_client.connect(self._cancel_event)
            self._log.info(
                "Yuanbao gateway connected: account=%s", self._account.account_id
            )
        except Exception:
            self._log.exception(
                "Failed to connect Yuanbao gateway: account=%s",
                self._account.account_id,
            )
            self._token_cache.invalidate(self._account.account_id)
            raise

        await self._sync_commands(access_token)

    async def _sync_commands(self, access_token: str) -> None:
        from yuxi.channel.extensions.yuanbao.commands import sync_native_commands

        await sync_native_commands(
            api_domain=self._account.api_domain,
            access_token=access_token,
        )

    async def _reconnect_loop(self) -> None:
        while not self._cancel_event.is_set():
            if self.is_connected:
                await asyncio.sleep(2)
                continue

            self._reconnect_count += 1
            if self._reconnect_count > self.MAX_RECONNECT_ATTEMPTS:
                self._log.error(
                    "Yuanbao gateway max reconnect attempts (%d) exceeded: account=%s",
                    self.MAX_RECONNECT_ATTEMPTS,
                    self._account.account_id,
                )
                break

            backoff = min(2 ** self._reconnect_count, 60)
            self._log.warning(
                "Yuanbao gateway reconnect attempt %d/%d in %ds: account=%s",
                self._reconnect_count,
                self.MAX_RECONNECT_ATTEMPTS,
                backoff,
                self._account.account_id,
            )

            await asyncio.sleep(backoff)
            if self._cancel_event.is_set():
                break

            try:
                self._token_cache.invalidate(self._account.account_id)
                await self._connect()
                self._reconnect_count = 0
            except Exception:
                self._log.exception(
                    "Yuanbao gateway reconnect failed: account=%s",
                    self._account.account_id,
                )

    async def stop(self) -> None:
        if self._listen_task:
            self._listen_task.cancel()
            self._listen_task = None
        if self._ws_client:
            await self._ws_client.disconnect()
            self._ws_client = None
        self._log.info(
            "Yuanbao gateway stopped: account=%s", self._account.account_id
        )

    async def send_text(
        self,
        target_id: str,
        content: str,
        reply_to_id: str | None = None,
    ) -> None:
        if not self._ws_client or not self._ws_client.is_connected:
            raise ConnectionError("Not connected")

        max_chars = self._config.max_chars if self._config else 3000
        if len(content) > max_chars:
            content = content[:max_chars]

        from yuxi.channel.extensions.yuanbao.codec.biz_codec import build_outbound_text

        msg = build_outbound_text(target_id, content, reply_to_id=reply_to_id)
        await self._ws_client.send_text_msg(msg)

    async def send_media(
        self,
        target_id: str,
        media_url: str,
        media_type: str,
        reply_to_id: str | None = None,
        file_name: str = "",
    ) -> None:
        if not self._ws_client or not self._ws_client.is_connected:
            raise ConnectionError("Not connected")

        from yuxi.channel.extensions.yuanbao.codec.biz_codec import (
            build_outbound_audio,
            build_outbound_file,
            build_outbound_image,
            build_outbound_video,
        )

        if media_type in ("image", "photo"):
            msg = build_outbound_image(target_id, media_url, reply_to_id=reply_to_id)
        elif media_type == "file":
            msg = build_outbound_file(target_id, media_url, file_name, reply_to_id=reply_to_id)
        elif media_type == "audio":
            msg = build_outbound_audio(target_id, media_url, reply_to_id=reply_to_id)
        elif media_type == "video":
            msg = build_outbound_video(target_id, media_url, reply_to_id=reply_to_id)
        else:
            msg = {
                "msgType": media_type,
                "targetId": target_id,
                "mediaUrl": media_url,
            }
            if reply_to_id:
                msg["replyToId"] = reply_to_id

        await self._ws_client.send_text_msg(msg)

    async def send_typing(self, target_id: str) -> None:
        if not self._ws_client or not self._ws_client.is_connected:
            return
        msg = {"msgType": "typing", "targetId": target_id}
        try:
            await self._ws_client.send_text_msg(msg)
        except Exception:
            self._log.debug("Failed to send typing indicator to %s", target_id)

    async def _on_ws_message(self, data: bytes) -> None:
        try:
            inbound = BizCodec.decode_inbound(data)
            unified = self._dispatcher.to_unified(
                inbound, self._account.account_id
            )
            if unified and self._queue is not None:
                await self._queue.put(unified)
        except Exception:
            self._log.exception("Error processing inbound Yuanbao message")