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

import asyncio
import logging
import random
import time

import httpx

from yuxi.channel.extensions.tencent_im.usersig import gen_user_sig

logger = logging.getLogger(__name__)

REGION_DOMAINS = {
    "ap-beijing": "console.tim.qq.com",
    "ap-guangzhou": "console.tim.qq.com",
    "ap-shanghai": "console.tim.qq.com",
    "ap-chengdu": "console.tim.qq.com",
    "ap-singapore": "adminapisgp.im.qcloud.com",
    "ap-seoul": "adminapikr.im.qcloud.com",
    "eu-frankfurt": "adminapiger.im.qcloud.com",
    "na-siliconvalley": "adminapiusa.im.qcloud.com",
    "ap-tokyo": "adminapijpn.im.qcloud.com",
    "ap-jakarta": "adminapijkt.im.qcloud.com",
}

RETRY_MAX = 3
RETRY_BASE_DELAY = 2
PROBE_INTERVAL = 60


class TencentIMError(Exception):
    def __init__(self, code: int, info: str = "", raw: dict | None = None):
        super().__init__(f"IM Error {code}: {info}")
        self.code = code
        self.info = info
        self.raw = raw or {}


class TencentIMGateway:
    def __init__(self):
        self._account = None
        self._admin_usersig: str | None = None
        self._usersig_expires_at: float = 0
        self._refresh_task: asyncio.Task | None = None
        self._probe_task: asyncio.Task | None = None
        self._http: httpx.AsyncClient | None = None
        self._running = False
        self._usersig_lock = asyncio.Lock()

    async def start(self, ctx) -> object:
        account = self._resolve_account(ctx)
        if not account.is_configured():
            logger.warning("TencentIM account not configured")
            return {"running": False, "reason": "not-configured"}

        self._account = account
        domain = REGION_DOMAINS.get(account.region, "console.tim.qq.com")
        self._http = httpx.AsyncClient(
            base_url=f"https://{domain}",
            timeout=account.http_timeout_ms / 1000,
        )

        try:
            await self._gen_admin_usersig()
            self._refresh_task = asyncio.create_task(self._usersig_refresh_loop())
            self._probe_task = asyncio.create_task(self._probe_loop())
            self._running = True
            logger.info("TencentIM gateway started for account %s", account.account_id)
            return {"running": True, "account_id": account.account_id}
        except Exception:
            logger.exception("TencentIM gateway failed to start")
            await self.stop(ctx)
            raise

    async def stop(self, ctx) -> None:
        self._running = False
        for task in [self._refresh_task, self._probe_task]:
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        if self._http:
            await self._http.aclose()
            self._http = None
        logger.info("TencentIM gateway stopped")

    async def _gen_admin_usersig(self):
        async with self._usersig_lock:
            self._admin_usersig = gen_user_sig(
                self._account.sdk_appid,
                self._account.secret_key,
                self._account.admin_userid,
                expire=self._account.usersig_expire_days * 86400,
            )
            self._usersig_expires_at = time.time() + self._account.usersig_expire_days * 86400

    def _build_request_url(self, service_cmd: str) -> str:
        return (
            f"/v4/{service_cmd}"
            f"?sdkappid={self._account.sdk_appid}"
            f"&identifier={self._account.admin_userid}"
            f"&usersig={self._admin_usersig}"
            f"&random={random.randint(0, 0xFFFFFFFF)}"
            f"&contenttype=json"
        )

    async def call_api(self, service_cmd: str, body: dict, retries: int = RETRY_MAX) -> dict:
        url = self._build_request_url(service_cmd)

        for attempt in range(retries):
            try:
                resp = await self._http.post(url, json=body)
                result = resp.json()
                error_code = result.get("ErrorCode", 0)

                if error_code == 0:
                    return result

                if error_code in (10003,):
                    delay = RETRY_BASE_DELAY * (2**attempt)
                    await asyncio.sleep(delay)
                    continue

                if error_code in (10004, 70001):
                    await self._gen_admin_usersig()
                    delay = RETRY_BASE_DELAY * (2**attempt)
                    await asyncio.sleep(delay)
                    continue

                raise TencentIMError(
                    code=error_code,
                    info=result.get("ErrorInfo", "unknown"),
                    raw=result,
                )

            except TencentIMError:
                raise
            except httpx.HTTPError as e:
                if attempt == retries - 1:
                    raise TencentIMError(code=-1, info=str(e))
                await asyncio.sleep(RETRY_BASE_DELAY * (2**attempt))

        raise TencentIMError(code=-1, info="max retries exceeded")

    async def _usersig_refresh_loop(self):
        refresh_interval = self._account.usersig_expire_days * 86400 / 2
        while self._running:
            try:
                await asyncio.sleep(refresh_interval)
                if self._running:
                    await self._gen_admin_usersig()
                    logger.info("TencentIM admin usersig refreshed")
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("TencentIM usersig refresh failed")

    async def _probe_loop(self):
        while self._running:
            try:
                await asyncio.sleep(PROBE_INTERVAL)
                if self._running:
                    await self.call_api(
                        "im_open_login_svc/account_check", {"CheckItem": [{"UserID": self._account.admin_userid}]}
                    )
            except asyncio.CancelledError:
                break
            except Exception:
                logger.warning("TencentIM probe failed")

    def get_gateway(self) -> "TencentIMGateway | None":
        return self if self._running else None

    @property
    def account(self):
        return self._account

    @staticmethod
    def _resolve_account(ctx):
        from yuxi.channel.extensions.tencent_im.config import TencentIMConfig

        config = TencentIMConfig()
        return config.resolve_account()