worldquant-alpha-system/backend/app/worldquant.py

"""Read-only WorldQuant adapter. Authentication is the only allowed upstream POST.

No upstream response body or request headers are included in exceptions: they may
contain credentials, cookies, or temporary authentication links.
"""

import asyncio
import math
import random
from datetime import datetime, timedelta, timezone
from email.utils import parsedate_to_datetime
from typing import Awaitable, Callable
from urllib.parse import urljoin, urlparse

import httpx


class WqError(Exception):
    def __init__(self, message: str, code: str = "upstream_error"):
        super().__init__(message)
        self.code = code


class VerificationRequired(WqError):
    def __init__(self, url: str):
        super().__init__("请在 WorldQuant 完成人工验证，再点击继续验证", "verification_required")
        self.url = url


class WqClient:
    def __init__(self, settings, transport=None, sleep=asyncio.sleep):
        self.settings = settings
        self.client = httpx.AsyncClient(
            base_url=settings.wq_base_url,
            timeout=settings.request_timeout,
            transport=transport,
            follow_redirects=False,
        )
        self.lock = asyncio.Lock()
        self.authenticated = False
        self.auth_generation = 0
        self.credentials: tuple[str, str] | None = None
        self.verification_url: str | None = None
        self.permissions: list[str] | None = None
        self.session_expires_at: datetime | None = None
        self.session_duration: float | None = None
        self.sleep = sleep
        self.on_retry: Callable[[float], Awaitable[None]] | None = None

    async def close(self):
        await self.client.aclose()

    def disconnect(self):
        self.authenticated = False
        self.credentials = None
        self.verification_url = None
        self.permissions = None
        self.session_expires_at = None
        self.session_duration = None
        self.client.cookies.clear()

    def capture_session(self, response: httpx.Response):
        """Keep only permissions and expiry metadata; never retain the authentication body."""
        self.permissions, self.session_expires_at, self.session_duration = None, None, None
        try:
            data = response.json()
        except ValueError:
            return
        if not isinstance(data, dict):
            return
        permissions = data.get("permissions")
        if isinstance(permissions, list) and all(isinstance(p, str) for p in permissions):
            self.permissions = list(dict.fromkeys(permissions))
        token = data.get("token")
        expiry = token.get("expiry") if isinstance(token, dict) else None
        if (
            isinstance(expiry, (int, float))
            and not isinstance(expiry, bool)
            and math.isfinite(expiry)
            and 0 <= expiry <= 31536000
        ):
            self.session_duration = expiry
            self.session_expires_at = datetime.now(timezone.utc) + timedelta(seconds=expiry)

    def session_info(self):
        remaining = (
            max(0, int((self.session_expires_at - datetime.now(timezone.utc)).total_seconds()))
            if self.session_expires_at
            else None
        )
        return {
            "authenticated": self.authenticated and remaining != 0,
            "expires_at": self.session_expires_at,
            "remaining_seconds": remaining,
            "total_seconds": self.session_duration,
        }

    def safe_verification_url(self, response: httpx.Response) -> str:
        url = urljoin(str(response.url), response.headers.get("Location", ""))
        base = urlparse(self.settings.wq_base_url)
        parsed = urlparse(url)
        if (
            not response.headers.get("Location")
            or parsed.scheme != base.scheme
            or parsed.netloc != base.netloc
        ):
            raise WqError("平台返回了无法确认的验证地址，请到 BRAIN 官网检查账户", "invalid_verification")
        return url

    @staticmethod
    def retry_delay(value: str | None, attempt: int) -> float:
        if value:
            try:
                seconds = float(value)
            except ValueError:
                try:
                    seconds = (parsedate_to_datetime(value) - datetime.now(timezone.utc)).total_seconds()
                except (TypeError, ValueError):
                    seconds = -1
            if math.isfinite(seconds) and seconds >= 0:
                return seconds
        return min(2**attempt + random.uniform(0, 0.5), 30)

    async def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
        for attempt in range(self.settings.retry_attempts):
            try:
                response = await self.client.request(method, path, **kwargs)
            except (httpx.TimeoutException, httpx.NetworkError):
                if attempt + 1 == self.settings.retry_attempts:
                    raise WqError("WorldQuant 网络请求失败，请稍后重试", "network_error") from None
                response = None
            if response is not None and response.status_code not in (429, 500, 502, 503, 504):
                return response
            if attempt + 1 == self.settings.retry_attempts:
                raise WqError("WorldQuant 暂时限流或不可用，重试预算已用完", "retry_exhausted")
            delay = self.retry_delay(
                response.headers.get("Retry-After") if response is not None else None, attempt
            )
            if self.on_retry:
                await self.on_retry(delay)
            await self.sleep(delay)
        raise WqError("重试失败")

    async def authenticate(self, email: str, password: str, force=False, stale_generation=None):
        async with self.lock:
            if (
                self.authenticated
                and not force
                and (stale_generation is None or stale_generation != self.auth_generation)
            ):
                return
            if self.verification_url and not force:
                raise VerificationRequired(self.verification_url)
            self.credentials = (email, password)
            self.authenticated = False
            self.permissions, self.session_expires_at, self.session_duration = None, None, None
            if force:
                self.client.cookies.clear()
                self.verification_url = None
            response = await self._request("POST", "/authentication", auth=httpx.BasicAuth(email, password))
            if (
                response.status_code == 401
                and response.headers.get("WWW-Authenticate", "").lower() == "persona"
            ):
                self.verification_url = self.safe_verification_url(response)
                raise VerificationRequired(self.verification_url)
            if response.status_code in (401, 403):
                raise WqError("平台认证未通过，请检查凭据或账户访问权限", "authentication_failed")
            if response.status_code not in (200, 201):
                raise WqError(f"平台认证返回异常状态 {response.status_code}", "authentication_failed")
            self.authenticated = True
            self.capture_session(response)
            self.auth_generation += 1
            self.verification_url = None

    async def verify(self):
        """Continue the same server-side challenge session after the human visits the URL."""
        if not self.verification_url:
            if not self.credentials:
                raise WqError("验证会话已结束，请重新连接", "authentication_failed")
            await self.authenticate(*self.credentials, force=True)
            return
        response = await self._request("POST", self.verification_url)
        if response.status_code in (200, 201):
            self.authenticated = True
            self.capture_session(response)
            self.auth_generation += 1
            self.verification_url = None
        elif response.status_code in (401, 403, 202):
            raise VerificationRequired(self.verification_url)
        else:
            raise WqError("验证会话已失效，请重新连接", "authentication_failed")

    async def get(self, path: str, params=None, headers=None):
        if not self.credentials:
            raise WqError("请先连接 WorldQuant", "disconnected")
        if not self.authenticated:
            await self.authenticate(*self.credentials)
        refreshed = False
        for attempt in range(self.settings.retry_attempts):
            generation = self.auth_generation
            response = await self._request("GET", path, params=params, headers=headers)
            if response.status_code == 401 and not refreshed:
                await self.authenticate(*self.credentials, stale_generation=generation)
                refreshed = True
                continue
            if response.status_code in (401, 403):
                raise WqError("无权访问该平台资源", "access_denied")
            if response.status_code == 404:
                raise WqError("平台资源不存在或不可访问", "not_found")
            if response.status_code not in (200, 201, 202):
                raise WqError(f"WorldQuant 请求返回状态 {response.status_code}")
            # Recordsets may return 200/202 with Retry-After before results exist.
            if (
                response.headers.get("Retry-After")
                and self.retry_delay(response.headers["Retry-After"], 0) > 0
            ):
                if attempt + 1 == self.settings.retry_attempts:
                    raise WqError("平台数据仍在准备，请稍后重试", "pending")
                delay = self.retry_delay(response.headers["Retry-After"], attempt)
                if self.on_retry:
                    await self.on_retry(delay)
                await self.sleep(delay)
                continue
            try:
                result = response.json()
                if not isinstance(result, dict):
                    raise ValueError()
                return result
            except ValueError:
                raise WqError("平台返回的数据格式无法识别", "invalid_response") from None
        raise WqError("平台数据尚未就绪", "pending")

    async def profile(self):
        return await self.get("/users/self")

    async def account_usage(self):
        """Read independent account resources; unavailable sections do not hide the profile."""
        from .account_data import usage_snapshot

        resources = {
            "simulations": "/users/self/activities/simulations",
            "submissions": "/users/self/activities/submissions",
            "alphas": "/users/self/alphas/summary",
        }
        data, errors = {}, {}
        for key, path in resources.items():
            try:
                data[key] = await self.get(
                    path, headers={"Accept": "application/json;version=4.0"} if key == "alphas" else None
                )
            except VerificationRequired:
                raise
            except WqError as exc:
                if exc.code in ("authentication_failed", "disconnected"):
                    raise
                errors[key] = str(exc)
        return usage_snapshot(data, errors)

    async def alphas(self, submission, hidden, offset, before):
        # Cover every platform stage; submitted records are not assumed to be OS only.
        status_key = "status" if submission == "UNSUBMITTED" else "status!"
        return await self.get(
            "/users/self/alphas",
            {
                status_key: "UNSUBMITTED",
                "hidden": str(hidden).lower(),
                "limit": 100,
                "offset": offset,
                "order": "dateCreated",
                "dateCreated<": before,
            },
        )

    async def alpha(self, alpha_id):
        return await self.get(f"/alphas/{alpha_id}")

    async def pnl(self, alpha_id):
        return await self.get(f"/alphas/{alpha_id}/recordsets/pnl")

    async def catalog_page(self, scope, dataset_id, offset):
        """Read a single scoped page. IDs are query parameters, never upstream paths."""
        params = {"instrumentType": scope["instrument_type"], "region": scope["region"],
                  "universe": scope["universe"], "delay": scope["delay"],
                  "limit": 50, "offset": offset}
        if dataset_id is not None:
            params["dataset.id"] = dataset_id
        return await self.get("/data-fields" if dataset_id else "/data-sets", params)