diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index 9197f85..12d31d6 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -9,7 +9,6 @@ import os import re import time from contextlib import asynccontextmanager -from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Optional, Tuple from urllib.parse import urlparse @@ -26,15 +25,15 @@ from app.core.config import get_config from app.core.exceptions import AppException, UpstreamException, ValidationException from app.core.logger import logger from app.core.storage import DATA_DIR -from app.services.grok.utils.headers import apply_statsig, build_sso_cookie -from app.services.token.service import TokenService +from app.services.reverse import ( + AssetsDeleteReverse, + AssetsDownloadReverse, + AssetsListReverse, + AssetsUploadReverse, +) # ==================== 常量 ==================== -UPLOAD_API = "https://grok.com/rest/app-chat/upload-file" -LIST_API = "https://grok.com/rest/assets" -DELETE_API = "https://grok.com/rest/assets-metadata" -DOWNLOAD_API = "https://assets.grok.com" LOCK_DIR = DATA_DIR / ".locks" # 全局信号量（运行时动态初始化） @@ -115,67 +114,15 @@ async def _file_lock(name: str, timeout: int = 10): fd.close() -@dataclass -class ServiceConfig: - """服务配置""" - - proxy: str - timeout: int - browser: str - user_agent: str - - @classmethod - def from_settings(cls, proxy: Optional[str] = None): - return cls( - proxy=proxy - or get_config("network.asset_proxy_url") - or get_config("network.base_proxy_url"), - timeout=get_config("network.timeout"), - browser=get_config("security.browser"), - user_agent=get_config("security.user_agent"), - ) - - def get_proxies(self) -> Optional[dict]: - return {"http": self.proxy, "https": self.proxy} if self.proxy else None - - # ==================== 基础服务 ==================== class BaseService: """基础服务类""" - def __init__(self, proxy: Optional[str] = None): - self.config = ServiceConfig.from_settings(proxy) + def __init__(self): self._session: Optional[AsyncSession] = None - def _build_headers( - self, token: str, referer: str = "https://grok.com/", download: bool = False - ) -> dict: - """构建请求头""" - if download: - headers = { - "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", - "Sec-Fetch-Dest": "document", - "Sec-Fetch-Mode": "navigate", - "Sec-Fetch-Site": "same-site", - "Sec-Fetch-User": "?1", - "Referer": referer, - "User-Agent": self.config.user_agent, - } - else: - headers = { - "Accept": "*/*", - "Content-Type": "application/json", - "Origin": "https://grok.com", - "Referer": referer, - "User-Agent": self.config.user_agent, - } - apply_statsig(headers) - - headers["Cookie"] = build_sso_cookie(token) - return headers - async def _get_session(self) -> AsyncSession: """获取复用 Session""" if self._session is None: @@ -298,44 +245,19 @@ class UploadService(BaseService): # 执行上传 session = await self._get_session() - response = await session.post( - UPLOAD_API, - headers=self._build_headers(token), - json={"fileName": filename, "fileMimeType": mime, "content": b64}, - impersonate=self.config.browser, - timeout=self.config.timeout, - proxies=self.config.get_proxies(), + response = await AssetsUploadReverse.request( + session, + token, + filename, + mime, + b64, ) - # 处理响应 - if response.status_code == 200: - result = response.json() - file_id = result.get("fileMetadataId", "") - file_uri = result.get("fileUri", "") - logger.info(f"Upload success: {filename} -> {file_id}") - return file_id, file_uri - - # 认证失败 - if response.status_code in (401, 403): - logger.warning(f"Upload auth failed: {response.status_code}") - try: - await TokenService.record_fail( - token, response.status_code, "upload_auth_failed" - ) - except Exception as e: - logger.error(f"Failed to record token failure: {e}") - - raise UpstreamException( - message=f"Upload authentication failed: {response.status_code}", - details={"status": response.status_code, "token_invalidated": True}, - ) - - # 其他错误 - logger.error(f"Upload failed: {filename} - {response.status_code}") - raise UpstreamException( - message=f"Upload failed: {response.status_code}", - details={"status": response.status_code}, - ) + result = response.json() + file_id = result.get("fileMetadataId", "") + file_uri = result.get("fileUri", "") + logger.info(f"Upload success: {filename} -> {file_id}") + return file_id, file_uri # ==================== 列表服务 ==================== @@ -346,7 +268,6 @@ class ListService(BaseService): async def iter_assets(self, token: str): """分页迭代资产列表""" - headers = self._build_headers(token, referer="https://grok.com/files") params = { "pageSize": 50, "orderBy": "ORDER_BY_LAST_USE_TIME", @@ -367,21 +288,12 @@ class ListService(BaseService): else: params.pop("pageToken", None) - response = await session.get( - LIST_API, - headers=headers, - params=params, - impersonate=self.config.browser, - timeout=self.config.timeout, - proxies=self.config.get_proxies(), + response = await AssetsListReverse.request( + session, + token, + params, ) - if response.status_code != 200: - raise UpstreamException( - message=f"List failed: {response.status_code}", - details={"status": response.status_code}, - ) - result = response.json() page_assets = result.get("assets", []) yield page_assets @@ -417,28 +329,19 @@ class DeleteService(BaseService): """删除单个文件""" async with _get_assets_semaphore(): session = await self._get_session() - response = await session.delete( - f"{DELETE_API}/{asset_id}", - headers=self._build_headers(token, referer="https://grok.com/files"), - impersonate=self.config.browser, - timeout=self.config.timeout, - proxies=self.config.get_proxies(), + response = await AssetsDeleteReverse.request( + session, + token, + asset_id, ) - if response.status_code == 200: - logger.debug(f"Deleted: {asset_id}") - return True - - logger.error(f"Delete failed: {asset_id} - {response.status_code}") - raise UpstreamException( - message=f"Delete failed: {asset_id}", - details={"status": response.status_code}, - ) + logger.debug(f"Deleted: {asset_id}") + return True async def delete_all(self, token: str) -> Dict[str, int]: """删除所有文件""" total = success = failed = 0 - list_service = ListService(self.config.proxy) + list_service = ListService() try: async for assets in list_service.iter_assets(token): @@ -496,8 +399,8 @@ class DeleteService(BaseService): class DownloadService(BaseService): """文件下载服务""" - def __init__(self, proxy: Optional[str] = None): - super().__init__(proxy) + def __init__(self): + super().__init__() self.base_dir = DATA_DIR / "tmp" self.image_dir = self.base_dir / "image" self.video_dir = self.base_dir / "video" @@ -552,25 +455,8 @@ class DownloadService(BaseService): if not file_path.startswith("/"): file_path = f"/{file_path}" - url = f"{DOWNLOAD_API}{file_path}" - headers = self._build_headers(token, download=True) - session = await self._get_session() - response = await session.get( - url, - headers=headers, - proxies=self.config.get_proxies(), - timeout=self.config.timeout, - allow_redirects=True, - impersonate=self.config.browser, - stream=True, - ) - - if response.status_code != 200: - raise UpstreamException( - message=f"Download failed: {response.status_code}", - details={"path": file_path, "status": response.status_code}, - ) + response = await AssetsDownloadReverse.request(session, token, file_path) # 保存文件 tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp") diff --git a/app/services/reverse/__init__.py b/app/services/reverse/__init__.py new file mode 100644 index 0000000..b86b1fd --- /dev/null +++ b/app/services/reverse/__init__.py @@ -0,0 +1,17 @@ +"""Reverse interfaces for Grok endpoints.""" + +from .assets_delete import AssetsDeleteReverse +from .assets_download import AssetsDownloadReverse +from .assets_list import AssetsListReverse +from .assets_upload import AssetsUploadReverse +from .utils.headers import build_headers +from .utils.statsig import StatsigGenerator + +__all__ = [ + "AssetsDeleteReverse", + "AssetsDownloadReverse", + "AssetsListReverse", + "AssetsUploadReverse", + "StatsigGenerator", + "build_headers", +] diff --git a/app/services/reverse/assets_delete.py b/app/services/reverse/assets_delete.py new file mode 100644 index 0000000..794f2ba --- /dev/null +++ b/app/services/reverse/assets_delete.py @@ -0,0 +1,102 @@ +""" +Reverse interface: delete asset metadata. +""" + +from typing import Any +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException +from app.services.token.service import TokenService +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status + +DELETE_API = "https://grok.com/rest/assets-metadata" + + +class AssetsDeleteReverse: + """/rest/assets-metadata/{file_id} reverse interface.""" + + @staticmethod + async def request(session: AsyncSession, token: str, asset_id: str) -> Any: + """Delete asset from Grok. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + asset_id: str, the ID of the asset to delete. + + Returns: + Any: The response from the request. + """ + try: + # Get proxies + base_proxy = get_config("network.base_proxy_url") + assert_proxy = get_config("network.asset_proxy_url") + if assert_proxy: + proxies = {"http": assert_proxy, "https": assert_proxy} + else: + proxies = {"http": base_proxy, "https": base_proxy} + + # Build headers + headers = build_headers( + cookie_token=token, + content_type="application/json", + origin="https://grok.com", + referer="https://grok.com/files", + ) + + # Curl Config + timeout = get_config("network.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.delete( + f"{DELETE_API}/{asset_id}", + headers=headers, + proxies=proxies, + timeout=timeout, + impersonate=browser, + ) + + if response.status_code != 200: + logger.error( + f"AssetsDeleteReverse: Delete failed, {response.status_code}", + extra={"error_type": "UpstreamException"}, + ) + raise UpstreamException( + message=f"AssetsDeleteReverse: Delete failed, {response.status_code}", + details={"status": response.status_code}, + ) + + return response + + return await retry_on_status(_do_request) + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + if status == 401: + try: + await TokenService.record_fail(token, status, "assets_delete_auth_failed") + except Exception: + pass + raise + + # Handle other non-upstream exceptions + logger.error( + f"AssetsDeleteReverse: Delete failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"AssetsDeleteReverse: Delete failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + +__all__ = ["AssetsDeleteReverse"] diff --git a/app/services/reverse/assets_download.py b/app/services/reverse/assets_download.py new file mode 100644 index 0000000..7b4addd --- /dev/null +++ b/app/services/reverse/assets_download.py @@ -0,0 +1,132 @@ +""" +Reverse interface: download asset. +""" + +import urllib.parse +from typing import Any +from pathlib import Path +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException +from app.services.token.service import TokenService +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status + +DOWNLOAD_API = "https://assets.grok.com" + +_CONTENT_TYPES = { + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".webp": "image/webp", + ".mp4": "video/mp4", + ".webm": "video/webm", +} + + +class AssetsDownloadReverse: + """assets.grok.com/{path} reverse interface.""" + + @staticmethod + async def request(session: AsyncSession, token: str, file_path: str) -> Any: + """Download asset from Grok. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + file_path: str, the path of the file to download. + + Returns: + Any: The response from the request. + """ + try: + # Normalize path + if not file_path.startswith("/"): + file_path = f"/{file_path}" + url = f"{DOWNLOAD_API}{file_path}" + + # Get proxies + base_proxy = get_config("network.base_proxy_url") + assert_proxy = get_config("network.asset_proxy_url") + if assert_proxy: + proxies = {"http": assert_proxy, "https": assert_proxy} + else: + proxies = {"http": base_proxy, "https": base_proxy} + + # Guess content type by extension for Accept/Sec-Fetch-Dest + content_type = _CONTENT_TYPES.get(Path(urllib.parse.urlparse(file_path).path).suffix.lower()) + + # Build headers + headers = build_headers( + cookie_token=token, + content_type=content_type, + origin="https://grok.com", + referer="https://grok.com/", + ) + ## Align with browser download navigation headers + headers["Cache-Control"] = "no-cache" + headers["Pragma"] = "no-cache" + headers["Priority"] = "u=0, i" + headers["Sec-Fetch-Mode"] = "navigate" + headers["Sec-Fetch-User"] = "?1" + headers["Upgrade-Insecure-Requests"] = "1" + + # Curl Config + timeout = get_config("network.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.get( + url, + headers=headers, + proxies=proxies, + timeout=timeout, + allow_redirects=True, + impersonate=browser, + stream=True, + ) + + if response.status_code != 200: + logger.error( + f"AssetsDownloadReverse: Download failed, {response.status_code}", + extra={"error_type": "UpstreamException"}, + ) + raise UpstreamException( + message=f"AssetsDownloadReverse: Download failed, {response.status_code}", + details={"status": response.status_code}, + ) + + return response + + return await retry_on_status(_do_request) + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + + if status == 401: + try: + await TokenService.record_fail(token, status, "assets_download_auth_failed") + except Exception: + pass + raise + + # Handle other non-upstream exceptions + logger.error( + f"AssetsDownloadReverse: Download failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"AssetsDownloadReverse: Download failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + + +__all__ = ["AssetsDownloadReverse"] diff --git a/app/services/reverse/assets_list.py b/app/services/reverse/assets_list.py new file mode 100644 index 0000000..0726372 --- /dev/null +++ b/app/services/reverse/assets_list.py @@ -0,0 +1,104 @@ +""" +Reverse interface: list assets. +""" + +from typing import Any, Dict +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException +from app.services.token.service import TokenService +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status + +LIST_API = "https://grok.com/rest/assets" + + +class AssetsListReverse: + """/rest/assets reverse interface.""" + + @staticmethod + async def request(session: AsyncSession, token: str, params: Dict[str, Any]) -> Any: + """List assets from Grok. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + params: Dict[str, Any], the parameters for the request. + + Returns: + Any: The response from the request. + """ + try: + # Get proxies + base_proxy = get_config("network.base_proxy_url") + assert_proxy = get_config("network.asset_proxy_url") + if assert_proxy: + proxies = {"http": assert_proxy, "https": assert_proxy} + else: + proxies = {"http": base_proxy, "https": base_proxy} + + # Build headers + headers = build_headers( + cookie_token=token, + content_type="application/json", + origin="https://grok.com", + referer="https://grok.com/files", + ) + + # Curl Config + timeout = get_config("network.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.get( + LIST_API, + headers=headers, + params=params, + proxies=proxies, + timeout=timeout, + impersonate=browser, + ) + + if response.status_code != 200: + logger.error( + f"AssetsListReverse: List failed, {response.status_code}", + extra={"error_type": "UpstreamException"}, + ) + raise UpstreamException( + message=f"AssetsListReverse: List failed, {response.status_code}", + details={"status": response.status_code}, + ) + + return response + + return await retry_on_status(_do_request) + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + if status == 401: + try: + await TokenService.record_fail(token, status, "assets_list_auth_failed") + except Exception: + pass + raise + + # Handle other non-upstream exceptions + logger.error( + f"AssetsListReverse: List failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"AssetsListReverse: List failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + + +__all__ = ["AssetsListReverse"] diff --git a/app/services/reverse/assets_upload.py b/app/services/reverse/assets_upload.py new file mode 100644 index 0000000..e4f65ac --- /dev/null +++ b/app/services/reverse/assets_upload.py @@ -0,0 +1,111 @@ +""" +Reverse interface: upload asset. +""" + +from typing import Any +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.exceptions import UpstreamException +from app.services.token.service import TokenService +from app.core.config import get_config +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status + +UPLOAD_API = "https://grok.com/rest/app-chat/upload-file" + + +class AssetsUploadReverse: + """/rest/app-chat/upload-file reverse interface.""" + + @staticmethod + async def request(session: AsyncSession, token: str, fileName: str, fileMimeType: str, content: str) -> Any: + """Upload asset to Grok. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + fileName: str, the name of the file. + fileMimeType: str, the MIME type of the file. + content: str, the content of the file. + + Returns: + Any: The response from the request. + """ + try: + # Get proxies + base_proxy = get_config("network.base_proxy_url") + assert_proxy = get_config("network.asset_proxy_url") + if assert_proxy: + proxies = {"http": assert_proxy, "https": assert_proxy} + else: + proxies = {"http": base_proxy, "https": base_proxy} + + # Build headers + headers = build_headers( + cookie_token=token, + content_type="application/json", + origin="https://grok.com", + referer="https://grok.com/", + ) + + # Build payload + payload = { + "fileName": fileName, + "fileMimeType": fileMimeType, + "content": content, + } + + # Curl Config + timeout = get_config("network.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.post( + UPLOAD_API, + headers=headers, + json=payload, + proxies=proxies, + timeout=timeout, + impersonate=browser, + ) + if response.status_code != 200: + logger.error( + f"AssetsUploadReverse: Upload failed, {response.status_code}", + extra={"error_type": "UpstreamException"}, + ) + raise UpstreamException( + message=f"AssetsUploadReverse: Upload failed, {response.status_code}", + details={"status": response.status_code}, + ) + return response + + return await retry_on_status(_do_request) + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + if status == 401: + try: + await TokenService.record_fail(token, status, "assets_upload_auth_failed") + except Exception: + pass + raise + + # Handle other non-upstream exceptions + logger.error( + f"AssetsUploadReverse: Upload failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"AssetsUploadReverse: Upload failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + + +__all__ = ["AssetsUploadReverse"] diff --git a/app/services/reverse/utils/headers.py b/app/services/reverse/utils/headers.py new file mode 100644 index 0000000..874425f --- /dev/null +++ b/app/services/reverse/utils/headers.py @@ -0,0 +1,101 @@ +"""Shared header builders for reverse interfaces.""" + +import uuid +import orjson +from typing import Dict, Optional + +from app.core.config import get_config +from app.core.logger import logger +from app.services.reverse.utils.statsig import StatsigGenerator + + +def _build_sso_cookie(sso_token: str) -> str: + """ + Build SSO Cookie string. + """ + # Remove "sso=" prefix if present + sso_token = sso_token[4:] if sso_token.startswith("sso=") else sso_token + + # SSO Cookie + cookie = f"sso={sso_token}; sso-rw={sso_token}" + + # CF Clearance + cf_clearance = get_config("security.cf_clearance") + if cf_clearance: + cookie += f";cf_clearance={cf_clearance}" + + return cookie + + +def build_headers( + cookie_token: str, + content_type: Optional[str] = None, + origin: Optional[str] = None, + referer: Optional[str] = None, +) -> Dict[str, str]: + """ + Build headers for reverse interfaces. + + Args: + cookie_token: The SSO token. + content_type: Optional Content-Type value. + origin: Optional Origin value. Defaults to "https://grok.com" if not provided. + referer: Optional Referer value. Defaults to "https://grok.com/" if not provided. + + Returns: + Dict[str, str]: The headers dictionary. + """ + headers = { + "Accept-Encoding": "gzip, deflate, br, zstd", + "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8", + "Baggage": "sentry-environment=production,sentry-release=d6add6fb0460641fd482d767a335ef72b9b6abb8,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c", + "Origin": origin or "https://grok.com", + "Priority": "u=1, i", + "Referer": referer or "https://grok.com/", + "Sec-Ch-Ua": '"Google Chrome";v="136", "Chromium";v="136", "Not(A:Brand";v="24"', + "Sec-Ch-Ua-Arch": "arm", + "Sec-Ch-Ua-Bitness": "64", + "Sec-Ch-Ua-Mobile": "?0", + "Sec-Ch-Ua-Model": "", + "Sec-Ch-Ua-Platform": '"macOS"', + "Sec-Fetch-Mode": "cors", + "User-Agent": get_config("security.user_agent"), + } + + # Cookie + headers["Cookie"] = _build_sso_cookie(cookie_token) + + # Content-Type and Accept/Sec-Fetch-Dest + if content_type and content_type == "application/json": + headers["Content-Type"] = "application/json" + headers["Accept"] = "*/*" + headers["Sec-Fetch-Dest"] = "empty" + elif content_type in ["image/jpeg", "image/png", "video/mp4", "video/webm"]: + headers["Content-Type"] = content_type + headers["Accept"] = "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7" + headers["Sec-Fetch-Dest"] = "document" + else: + headers["Content-Type"] = "application/json" + headers["Accept"] = "*/*" + headers["Sec-Fetch-Dest"] = "empty" + + # Sec-Fetch-Site + if headers["Origin"] == headers["Referer"]: + headers["Sec-Fetch-Site"] = "same-origin" + else: + headers["Sec-Fetch-Site"] = "same-site" + + # X-Statsig-ID and X-XAI-Request-ID + headers["x-statsig-id"] = StatsigGenerator.gen_id() + headers["x-xai-request-id"] = str(uuid.uuid4()) + + # Print headers without Cookie + safe_headers = dict(headers) + if "Cookie" in safe_headers: + safe_headers["Cookie"] = "" + logger.debug(f"Built headers: {orjson.dumps(safe_headers, indent=2)}") + + return headers + + +__all__ = ["build_headers"] diff --git a/app/services/reverse/utils/retry.py b/app/services/reverse/utils/retry.py new file mode 100644 index 0000000..0de15b6 --- /dev/null +++ b/app/services/reverse/utils/retry.py @@ -0,0 +1,229 @@ +""" +Reverse retry utilities. +""" + +import asyncio +import random +from typing import Callable, Any, Optional + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException + + +class RetryContext: + """Retry context.""" + + def __init__(self): + self.attempt = 0 + self.max_retry = int(get_config("retry.max_retry")) + self.retry_codes = get_config("retry.retry_status_codes") + self.last_error = None + self.last_status = None + self.total_delay = 0.0 + self.retry_budget = float(get_config("retry.retry_budget")) + + # Backoff parameters + self.backoff_base = float(get_config("retry.retry_backoff_base")) + self.backoff_factor = float(get_config("retry.retry_backoff_factor")) + self.backoff_max = float(get_config("retry.retry_backoff_max")) + + # Decorrelated jitter state + self._last_delay = self.backoff_base + + def should_retry(self, status_code: int) -> bool: + """Check if should retry.""" + if self.attempt >= self.max_retry: + return False + if status_code not in self.retry_codes: + return False + if self.total_delay >= self.retry_budget: + return False + return True + + def record_error(self, status_code: int, error: Exception): + """Record error information.""" + self.last_status = status_code + self.last_error = error + self.attempt += 1 + + def calculate_delay(self, status_code: int, retry_after: Optional[float] = None) -> float: + """ + Calculate backoff delay time. + + Args: + status_code: HTTP status code + retry_after: Retry-After header value (seconds) + + Returns: + Delay time (seconds) + """ + # Use Retry-After if available + if retry_after is not None and retry_after > 0: + delay = min(retry_after, self.backoff_max) + self._last_delay = delay + return delay + + # Use decorrelated jitter for 429 + if status_code == 429: + # decorrelated jitter: delay = random(base, last_delay * 3) + delay = random.uniform(self.backoff_base, self._last_delay * 3) + delay = min(delay, self.backoff_max) + self._last_delay = delay + return delay + + # Use exponential backoff + full jitter for other status codes + exp_delay = self.backoff_base * (self.backoff_factor**self.attempt) + delay = random.uniform(0, min(exp_delay, self.backoff_max)) + return delay + + def record_delay(self, delay: float): + """Record delay time.""" + self.total_delay += delay + + +def extract_retry_after(error: Exception) -> Optional[float]: + """ + Extract Retry-After value from exception. + + Args: + error: Exception object + + Returns: + Retry-After value (seconds), or None + """ + if not isinstance(error, UpstreamException): + return None + + details = error.details or {} + + # Try to get Retry-After from details + retry_after = details.get("retry_after") + if retry_after is not None: + try: + return float(retry_after) + except (ValueError, TypeError): + pass + + # Try to get Retry-After from headers + headers = details.get("headers", {}) + if isinstance(headers, dict): + retry_after = headers.get("Retry-After") or headers.get("retry-after") + if retry_after is not None: + try: + return float(retry_after) + except (ValueError, TypeError): + pass + + return None + + +async def retry_on_status( + func: Callable, + *args, + extract_status: Callable[[Exception], Optional[int]] = None, + on_retry: Callable[[int, int, Exception, float], None] = None, + **kwargs, +) -> Any: + """ + Generic retry function. + + Args: + func: Retry function + *args: Function arguments + extract_status: Function to extract status code from exception + on_retry: Callback function for retry (attempt, status_code, error, delay) + **kwargs: Function keyword arguments + + Returns: + Function execution result + + Raises: + Last failed exception + """ + ctx = RetryContext() + + # Status code extractor + if extract_status is None: + + def extract_status(e: Exception) -> Optional[int]: + if isinstance(e, UpstreamException): + # Try to get status code from details, fallback to status_code attribute + if e.details and "status" in e.details: + return e.details["status"] + return getattr(e, "status_code", None) + return None + + while ctx.attempt <= ctx.max_retry: + try: + result = await func(*args, **kwargs) + + # Record log + if ctx.attempt > 0: + logger.info( + f"Retry succeeded after {ctx.attempt} attempts, " + f"total delay: {ctx.total_delay:.2f}s" + ) + + return result + + except Exception as e: + # Extract status code + status_code = extract_status(e) + + if status_code is None: + # Error cannot be identified as retryable + logger.error(f"Non-retryable error: {e}") + raise + + # Record error + ctx.record_error(status_code, e) + + # Check if should retry + if ctx.should_retry(status_code): + # Extract Retry-After + retry_after = extract_retry_after(e) + + # Calculate delay + delay = ctx.calculate_delay(status_code, retry_after) + + # Check if exceeds budget + if ctx.total_delay + delay > ctx.retry_budget: + logger.warning( + f"Retry budget exhausted: {ctx.total_delay:.2f}s + {delay:.2f}s > {ctx.retry_budget}s" + ) + raise + + ctx.record_delay(delay) + + logger.warning( + f"Retry {ctx.attempt}/{ctx.max_retry} for status {status_code}, " + f"waiting {delay:.2f}s (total: {ctx.total_delay:.2f}s)" + + (f", Retry-After: {retry_after}s" if retry_after else "") + ) + + # Callback + if on_retry: + on_retry(ctx.attempt, status_code, e, delay) + + await asyncio.sleep(delay) + continue + else: + # Not retryable or retry budget exhausted + if status_code in ctx.retry_codes: + logger.error( + f"Retry exhausted after {ctx.attempt} attempts, " + f"last status: {status_code}, total delay: {ctx.total_delay:.2f}s" + ) + else: + logger.error(f"Non-retryable status code: {status_code}") + + # Raise last failed exception + raise + + +__all__ = [ + "RetryContext", + "retry_on_status", + "extract_retry_after", +] diff --git a/app/services/reverse/utils/statsig.py b/app/services/reverse/utils/statsig.py new file mode 100644 index 0000000..28280a0 --- /dev/null +++ b/app/services/reverse/utils/statsig.py @@ -0,0 +1,59 @@ +""" +Statsig ID generator for reverse interfaces. +""" + +import base64 +import random +import string + +from app.core.config import get_config +from app.core.logger import logger + + +STATIC_STATSIG_ID = "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" + + +class StatsigGenerator: + """Statsig ID generator for reverse interfaces.""" + + @staticmethod + def _rand(length: int, alphanumeric: bool = False) -> str: + """Generate random string.""" + chars = ( + string.ascii_lowercase + string.digits + if alphanumeric + else string.ascii_lowercase + ) + return "".join(random.choices(chars, k=length)) + + @staticmethod + def gen_id() -> str: + """ + Generate Statsig ID. + + Returns: + Base64 encoded ID. + """ + dynamic = get_config("chat.dynamic_statsig") + + # Dynamic Statsig ID + if dynamic: + logger.debug("Generating dynamic Statsig ID") + + if random.choice([True, False]): + rand = StatsigGenerator._rand(5, alphanumeric=True) + message = f"e:TypeError: Cannot read properties of null (reading 'children['{rand}']')" + else: + rand = StatsigGenerator._rand(10) + message = ( + f"e:TypeError: Cannot read properties of undefined (reading '{rand}')" + ) + + return base64.b64encode(message.encode()).decode() + + # Static Statsig ID + logger.debug("Generating static Statsig ID") + return STATIC_STATSIG_ID + + +__all__ = ["StatsigGenerator"]