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

import asyncio
import logging
import time

import httpx

logger = logging.getLogger(__name__)

TOKEN_URL = "https://open.xiaohongshu.com/oauth/access_token"
REFRESH_TOKEN_URL = "https://open.xiaohongshu.com/oauth/refresh_token"
REFRESH_INTERVAL = 3600
RETRY_MAX = 3
RETRY_BASE_DELAY = 2

_gateway: "XiaohongshuGateway | None" = None


def _set_gateway(gw: "XiaohongshuGateway") -> None:
    global _gateway
    _gateway = gw


def _get_gateway() -> "XiaohongshuGateway | None":
    return _gateway


class XiaohongshuAuthError(Exception):
    pass


class XiaohongshuGateway:

    def __init__(self):
        self._account = None
        self._access_token: str | None = None
        self._refresh_token: str | None = None
        self._token_expires_at: float = 0
        self._refresh_task: asyncio.Task | None = None
        self._cancel_event = asyncio.Event()
        self._http: httpx.AsyncClient | None = None
        self._running = False
        self._token_lock = asyncio.Lock()

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

        self._account = account
        self._cancel_event.clear()
        self._http = httpx.AsyncClient(
            base_url="https://open.xiaohongshu.com",
            timeout=15.0,
            headers={"Content-Type": "application/json", "User-Agent": "ForcePilot-XHS/1.0"},
        )

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

    async def stop(self, ctx) -> None:
        self._running = False
        self._cancel_event.set()

        if self._refresh_task and not self._refresh_task.done():
            self._refresh_task.cancel()
            try:
                await self._refresh_task
            except asyncio.CancelledError:
                pass

        if self._http:
            await self._http.aclose()
            self._http = None

        logger.info("Xiaohongshu gateway stopped")

    @property
    def access_token(self) -> str | None:
        if self._access_token and time.time() < self._token_expires_at:
            return self._access_token
        return None

    async def _ensure_valid_token(self) -> None:
        if self._access_token and time.time() < self._token_expires_at:
            return
        if self._refresh_token:
            await self._refresh_access_token()
        else:
            raise XiaohongshuAuthError("No valid token and no refresh token available")

    async def _refresh_access_token(self) -> str:
        async with self._token_lock:
            if self._access_token and time.time() < self._token_expires_at:
                return self._access_token

            for attempt in range(RETRY_MAX):
                try:
                    resp = await self._http.post(
                        "/oauth/refresh_token",
                        json={
                            "app_key": self._account.app_key,
                            "refresh_token": self._refresh_token,
                            "grant_type": "refresh_token",
                        },
                    )
                    data = resp.json()
                    if data.get("code") != 0:
                        raise XiaohongshuAuthError(f"Failed to refresh token: {data}")

                    self._access_token = data["data"]["access_token"]
                    self._refresh_token = data["data"].get("refresh_token", self._refresh_token)
                    expires_in = data["data"].get("expires_in", 7200)
                    self._token_expires_at = time.time() + expires_in - 300
                    logger.info("Xiaohongshu access_token refreshed, expires_in=%ds", expires_in)
                    return self._access_token

                except Exception:
                    if attempt == RETRY_MAX - 1:
                        raise
                delay = RETRY_BASE_DELAY * (2 ** attempt)
                await asyncio.sleep(delay)

            raise XiaohongshuAuthError("Failed to refresh access_token after retries")

    async def _token_refresh_loop(self):
        while not self._cancel_event.is_set():
            try:
                await asyncio.sleep(REFRESH_INTERVAL)
                if not self._cancel_event.is_set():
                    await self._ensure_valid_token()
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("Xiaohongshu access_token refresh failed in loop")

    @staticmethod
    def _resolve_account(ctx):
        from yuxi.channel.extensions.xiaohongshu.config import XiaohongshuConfig

        config = XiaohongshuConfig()
        return config.resolve_account()