refactor: integrate reverse request handlers for image and voice services

app/services/grok/services/image.py

@@ -10,14 +10,13 @@ import ssl
 import time
 import uuid
 from typing import AsyncGenerator, Dict, Optional
-from urllib.parse import urlparse
 
 import aiohttp
-from aiohttp_socks import ProxyConnector
 
 from app.core.config import get_config
 from app.core.logger import logger
 from app.services.grok.utils.headers import build_sso_cookie
+from app.services.reverse.utils.websocket import resolve_proxy
 
 WS_URL = "wss://grok.com/ws/imagine/listen"
 
@@ -36,16 +35,7 @@ class ImageService:
 
     def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]:
         proxy_url = get_config("network.base_proxy_url")
-        if not proxy_url:
-            return aiohttp.TCPConnector(ssl=self._ssl_context), None
-
-        scheme = urlparse(proxy_url).scheme.lower()
-        if scheme.startswith("socks"):
-            logger.info(f"Using SOCKS proxy: {proxy_url}")
-            return ProxyConnector.from_url(proxy_url, ssl=self._ssl_context), None
-
-        logger.info(f"Using HTTP proxy: {proxy_url}")
-        return aiohttp.TCPConnector(ssl=self._ssl_context), proxy_url
+        return resolve_proxy(proxy_url, self._ssl_context)
 
     def _get_ws_headers(self, token: str) -> Dict[str, str]:
         cookie = build_sso_cookie(token, include_rw=True)

app/services/grok/services/voice.py

@@ -2,25 +2,17 @@
 Grok Voice Mode Service
 """
 
-import orjson
-from typing import Dict, Any
+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.grok.utils.headers import apply_statsig, build_sso_cookie
-
-LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens"
+from app.services.reverse import LivekitTokenReverse
 
 
 class VoiceService:
     """Voice Mode Service (LiveKit)"""
 
-    def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("network.base_proxy_url")
-
     async def get_token(
         self,
         token: str,
@@ -28,86 +20,12 @@ class VoiceService:
         personality: str = "assistant",
         speed: float = 1.0,
     ) -> Dict[str, Any]:
-        """
-        Get LiveKit token
-
-        Args:
-            token: Auth token
-        Returns:
-            Dict containing token and livekitUrl
-        """
-        logger.debug(
-            f"Voice token request: voice={voice}, personality={personality}, speed={speed}"
-        )
-        headers = self._build_headers(token)
-        payload = self._build_payload(voice, personality, speed)
-
-        proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None
-
-        try:
-            browser = get_config("security.browser")
-            timeout = get_config("network.timeout")
-            async with AsyncSession(impersonate=browser) as session:
-                response = await session.post(
-                    LIVEKIT_TOKEN_API,
-                    headers=headers,
-                    data=orjson.dumps(payload),
-                    timeout=timeout,
-                    proxies=proxies,
-                )
-
-                if response.status_code != 200:
-                    body = response.text[:200]
-                    logger.error(
-                        f"Voice token failed: status={response.status_code}, body={body}"
-                    )
-                    raise UpstreamException(
-                        message=f"Failed to get voice token: {response.status_code}",
-                        details={"status": response.status_code, "body": response.text},
-                    )
-
-                result = response.json()
-                logger.info(f"Voice token obtained: voice={voice}")
-                return result
-
-        except Exception as e:
-            logger.error(f"Voice service error: {e}")
-            if isinstance(e, UpstreamException):
-                raise
-            raise UpstreamException(f"Voice service error: {str(e)}")
-
-    def _build_headers(self, token: str) -> Dict[str, str]:
-        headers = {
-            "Accept": "*/*",
-            "Content-Type": "application/json",
-            "Origin": "https://grok.com",
-            "Referer": "https://grok.com/",
-            # Statsig ID is crucial
-        }
-
-        apply_statsig(headers)
-        headers["Cookie"] = build_sso_cookie(token)
-
-        return headers
-
-    def _build_payload(
-        self,
-        voice: str = "ara",
-        personality: str = "assistant",
-        speed: float = 1.0,
-    ) -> Dict[str, Any]:
-        """Construct payload with voice settings"""
-        return {
-            "sessionPayload": orjson.dumps(
-                {
-                    "voice": voice,
-                    "personality": personality,
-                    "playback_speed": speed,
-                    "enable_vision": False,
-                    "turn_detection": {"type": "server_vad"},
-                }
-            ).decode(),
-            "requestAgentDispatch": False,
-            "livekitUrl": "wss://livekit.grok.com",
-            "params": {"enable_markdown_transcript": "true"},
-        }
+        browser = get_config("security.browser")
+        async with AsyncSession(impersonate=browser) as session:
+            return await LivekitTokenReverse.request(
+                session,
+                token=token,
+                voice=voice,
+                personality=personality,
+                speed=speed,
+            )

app/services/reverse/__init__.py

@@ -9,8 +9,7 @@ from .media_post import MediaPostReverse
 from .nsfw_mgmt import NsfwMgmtReverse
 from .rate_limits import RateLimitsReverse
 from .set_birth import SetBirthReverse
-from .livekit_tokens import LivekitTokenReverse
-from .ws_livekit import LivekitWebSocketReverse
+from .ws_livekit import LivekitTokenReverse, LivekitWebSocketReverse
 from .utils.headers import build_headers
 from .utils.statsig import StatsigGenerator
 

app/services/reverse/assets_upload.py

@@ -6,9 +6,9 @@ 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.core.config import get_config
 from app.services.reverse.utils.headers import build_headers
 from app.services.reverse.utils.retry import retry_on_status
 

app/services/reverse/utils/grpc.py

@@ -10,9 +10,6 @@ from typing import Dict, List, Mapping, Tuple
 from urllib.parse import unquote
 
 
-_B64_RE = re.compile(rb"^[A-Za-z0-9+/=\r\n]+$")
-
-
 @dataclass(frozen=True)
 class GrpcStatus:
     code: int
@@ -51,7 +48,7 @@ class GrpcClient:
             return base64.b64decode(compact, validate=False)
 
         head = body[: min(len(body), 2048)]
-        if head and _B64_RE.fullmatch(head):
+        if head and re.compile(rb"^[A-Za-z0-9+/=\r\n]+$").fullmatch(head):
             compact = b"".join(body.split())
             try:
                 return base64.b64decode(compact, validate=True)

app/services/reverse/utils/headers.py

@@ -5,8 +5,8 @@ import orjson
 from urllib.parse import urlparse
 from typing import Dict, Optional
 
-from app.core.config import get_config
 from app.core.logger import logger
+from app.core.config import get_config
 from app.services.reverse.utils.statsig import StatsigGenerator
 
 
@@ -14,7 +14,7 @@ def _build_sso_cookie(sso_token: str) -> str:
     """
     Build SSO Cookie string.
     """
-    # Remove "sso=" prefix if present
+    # Format
     sso_token = sso_token[4:] if sso_token.startswith("sso=") else sso_token
 
     # SSO Cookie

app/services/reverse/utils/statsig.py

@@ -6,11 +6,8 @@ import base64
 import random
 import string
 
-from app.core.config import get_config
 from app.core.logger import logger
-
-
-STATIC_STATSIG_ID = "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk="
+from app.core.config import get_config
 
 
 class StatsigGenerator:
@@ -53,7 +50,7 @@ class StatsigGenerator:
 
         # Static Statsig ID
         logger.debug("Generating static Statsig ID")
-        return STATIC_STATSIG_ID
+        return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk="
 
 
 __all__ = ["StatsigGenerator"]

app/services/reverse/utils/websocket.py

@@ -0,0 +1,144 @@
+"""
+WebSocket helpers for reverse interfaces.
+"""
+
+import ssl
+import certifi
+import aiohttp
+from aiohttp_socks import ProxyConnector
+from typing import Mapping, Optional
+from urllib.parse import urlparse
+
+from app.core.logger import logger
+from app.core.config import get_config
+
+
+def _default_ssl_context() -> ssl.SSLContext:
+    context = ssl.create_default_context()
+    context.load_verify_locations(certifi.where())
+    return context
+
+
+def _normalize_socks_proxy(proxy_url: str) -> tuple[str, Optional[bool]]:
+    scheme = urlparse(proxy_url).scheme.lower()
+    rdns: Optional[bool] = None
+    base_scheme = scheme
+
+    if scheme == "socks5h":
+        base_scheme = "socks5"
+        rdns = True
+    elif scheme == "socks4a":
+        base_scheme = "socks4"
+        rdns = True
+
+    if base_scheme != scheme:
+        proxy_url = proxy_url.replace(f"{scheme}://", f"{base_scheme}://", 1)
+
+    return proxy_url, rdns
+
+
+def resolve_proxy(
+    proxy_url: str | None, ssl_context: ssl.SSLContext
+) -> tuple[aiohttp.BaseConnector, Optional[str]]:
+    """Resolve proxy connector.
+    
+    Args:
+        proxy_url: str, the proxy URL.
+        ssl_context: ssl.SSLContext, the SSL context.
+
+    Returns:
+        tuple[aiohttp.BaseConnector, Optional[str]]: The proxy connector and the proxy URL.
+    """
+    if not proxy_url:
+        return aiohttp.TCPConnector(ssl=ssl_context), None
+
+    scheme = urlparse(proxy_url).scheme.lower()
+    if scheme.startswith("socks"):
+        normalized, rdns = _normalize_socks_proxy(proxy_url)
+        logger.info(f"Using SOCKS proxy: {proxy_url}")
+        try:
+            if rdns is not None:
+                return (
+                    ProxyConnector.from_url(normalized, rdns=rdns, ssl=ssl_context),
+                    None,
+                )
+        except TypeError:
+            return ProxyConnector.from_url(normalized, ssl=ssl_context), None
+        return ProxyConnector.from_url(normalized, ssl=ssl_context), None
+
+    logger.info(f"Using HTTP proxy: {proxy_url}")
+    return aiohttp.TCPConnector(ssl=ssl_context), proxy_url
+
+
+class WebSocketConnection:
+    """WebSocket connection wrapper."""
+
+    def __init__(
+        self,
+        session: aiohttp.ClientSession,
+        ws: aiohttp.ClientWebSocketResponse,
+    ) -> None:
+        self.session = session
+        self.ws = ws
+
+    async def close(self) -> None:
+        if not self.ws.closed:
+            await self.ws.close()
+        await self.session.close()
+
+    async def __aenter__(self) -> aiohttp.ClientWebSocketResponse:
+        return self.ws
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self.close()
+
+
+class WebSocketClient:
+    """WebSocket client with proxy support."""
+
+    def __init__(self, proxy: str | None = None) -> None:
+        self.proxy = proxy or get_config("network.base_proxy_url")
+        self._ssl_context = _default_ssl_context()
+
+    async def connect(
+        self,
+        url: str,
+        headers: Mapping[str, str] | None = None,
+        timeout: float | aiohttp.ClientTimeout | None = None,
+    ) -> WebSocketConnection:
+        """Connect to the WebSocket.
+        
+        Args:
+            url: str, the URL to connect to.
+            headers: Mapping[str, str], the headers to send.
+            timeout: float | aiohttp.ClientTimeout | None, the timeout.
+
+        Returns:
+            WebSocketConnection: The WebSocket connection.
+        """
+        # Resolve proxy
+        connector, proxy = resolve_proxy(self.proxy, self._ssl_context)
+
+        # Build client timeout
+        client_timeout = (
+            timeout
+            if isinstance(timeout, aiohttp.ClientTimeout)
+            else aiohttp.ClientTimeout(total=timeout)
+        )
+
+        # Create session
+        session = aiohttp.ClientSession(connector=connector, timeout=client_timeout)
+        try:
+            ws = await session.ws_connect(
+                url,
+                headers=headers,
+                proxy=proxy,
+                ssl=self._ssl_context,
+            )
+            return WebSocketConnection(session, ws)
+        except Exception:
+            await session.close()
+            raise
+
+
+__all__ = ["WebSocketClient", "WebSocketConnection", "resolve_proxy"]

app/services/reverse/ws_livekit.py

@@ -0,0 +1,253 @@
+"""
+Reverse interface: LiveKit token + WebSocket.
+"""
+
+import orjson
+from typing import Any, Dict
+from urllib.parse import urlencode
+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
+from app.services.reverse.utils.websocket import WebSocketClient, WebSocketConnection
+
+LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens"
+LIVEKIT_WS_URL = "wss://livekit.grok.com"
+
+
+class LivekitTokenReverse:
+    """/rest/livekit/tokens reverse interface."""
+
+    @staticmethod
+    async def request(
+        session: AsyncSession,
+        token: str,
+        voice: str = "ara",
+        personality: str = "assistant",
+        speed: float = 1.0,
+        livekit_url: str = LIVEKIT_WS_URL,
+    ) -> Dict[str, Any]:
+        """Fetch LiveKit token.
+        
+        Args:
+            session: AsyncSession, the session to use for the request.
+            token: str, the SSO token.
+            voice: str, the voice to use for the request.
+            personality: str, the personality to use for the request.
+            speed: float, the speed to use for the request.
+            livekit_url: str, the LiveKit URL to use for the request.
+
+        Returns:
+            Dict[str, Any]: The LiveKit token.
+        """
+        try:
+            # Get proxies
+            base_proxy = get_config("network.base_proxy_url")
+            proxies = {"http": base_proxy, "https": base_proxy} if base_proxy else None
+
+            # Build headers
+            headers = build_headers(
+                cookie_token=token,
+                content_type="application/json",
+                origin="https://grok.com",
+                referer="https://grok.com/",
+            )
+
+            # Build payload
+            payload = {
+                "sessionPayload": orjson.dumps(
+                    {
+                        "voice": voice,
+                        "personality": personality,
+                        "playback_speed": speed,
+                        "enable_vision": False,
+                        "turn_detection": {"type": "server_vad"},
+                    }
+                ).decode(),
+                "requestAgentDispatch": False,
+                "livekitUrl": livekit_url,
+                "params": {"enable_markdown_transcript": "true"},
+            }
+
+            # Curl Config
+            timeout = get_config("network.timeout")
+            browser = get_config("security.browser")
+
+            async def _do_request():
+                response = await session.post(
+                    LIVEKIT_TOKEN_API,
+                    headers=headers,
+                    data=orjson.dumps(payload),
+                    timeout=timeout,
+                    proxies=proxies,
+                    impersonate=browser,
+                )
+
+                if response.status_code != 200:
+                    body = response.text[:200]
+                    logger.error(
+                        f"LivekitTokenReverse: Request failed, {response.status_code}, body={body}"
+                    )
+                    raise UpstreamException(
+                        message=f"LivekitTokenReverse: Request failed, {response.status_code}",
+                        details={"status": response.status_code, "body": response.text},
+                    )
+
+                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, "livekit_token_auth_failed"
+                        )
+                    except Exception:
+                        pass
+                raise
+
+            # Handle other non-upstream exceptions
+            logger.error(
+                f"LivekitTokenReverse: Request failed, {str(e)}",
+                extra={"error_type": type(e).__name__},
+            )
+            raise UpstreamException(
+                message=f"LivekitTokenReverse: Request failed, {str(e)}",
+                details={"status": 502, "error": str(e)},
+            )
+
+
+class LivekitWebSocketReverse:
+    """LiveKit WebSocket reverse interface."""
+
+    def __init__(self) -> None:
+        self._client = WebSocketClient()
+
+    def build_url(
+        self,
+        access_token: str,
+        *,
+        livekit_url: str = LIVEKIT_WS_URL,
+        auto_subscribe: bool = True,
+        sdk: str = "js",
+        version: str = "2.11.4",
+        protocol: int = 15,
+    ) -> str:
+        """Build LiveKit WebSocket URL.
+        
+        Args:
+            access_token: str, the LiveKit access token.
+            livekit_url: str, the LiveKit URL to use for the request.
+            auto_subscribe: bool, whether to auto subscribe to the WebSocket.
+            sdk: str, the SDK to use for the request.
+            version: str, the version to use for the request.
+            protocol: int, the protocol to use for the request.
+
+        Returns:
+            str: The LiveKit WebSocket URL.
+        """
+        # Build base URL
+        base = livekit_url.rstrip("/")
+        if not base.endswith("/rtc"):
+            base = f"{base}/rtc"
+
+        # Build parameters
+        params = {
+            "access_token": access_token,
+            "auto_subscribe": str(int(auto_subscribe)),
+            "sdk": sdk,
+            "version": version,
+            "protocol": str(protocol),
+        }
+
+        return f"{base}?{urlencode(params)}"
+
+    def _build_headers(self, extra: Dict[str, str] | None = None) -> Dict[str, str]:
+        """Build LiveKit WebSocket headers."""
+        # Build headers
+        headers = {
+            "Origin": "https://grok.com",
+            "User-Agent": get_config("security.user_agent"),
+            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
+            "Cache-Control": "no-cache",
+            "Pragma": "no-cache",
+        }
+
+        # Update headers
+        if extra:
+            headers.update(extra)
+        return headers
+
+    async def connect(
+        self,
+        access_token: str,
+        *,
+        livekit_url: str = LIVEKIT_WS_URL,
+        auto_subscribe: bool = True,
+        sdk: str = "js",
+        version: str = "2.11.4",
+        protocol: int = 15,
+        headers: Dict[str, str] | None = None,
+        timeout: float | None = None,
+    ) -> WebSocketConnection:
+        """Connect to the LiveKit WebSocket.
+        
+        Args:
+            access_token: str, the LiveKit access token.
+            livekit_url: str, the LiveKit URL to use for the request.
+            auto_subscribe: bool, whether to auto subscribe to the WebSocket.
+            sdk: str, the SDK to use for the request.
+            version: str, the version to use for the request.
+            protocol: int, the protocol to use for the request.
+            headers: Dict[str, str], the headers to send.
+            timeout: float, the timeout to use for the request.
+
+        Returns:
+            WebSocketConnection: The LiveKit WebSocket connection.
+        """
+        # Build URL
+        url = self.build_url(
+            access_token,
+            livekit_url=livekit_url,
+            auto_subscribe=auto_subscribe,
+            sdk=sdk,
+            version=version,
+            protocol=protocol,
+        )
+
+        # Build WebSocket headers
+        ws_headers = self._build_headers(headers)
+
+        # Build timeout
+        if timeout is None:
+            timeout = get_config("network.timeout")
+
+        # Connect to the LiveKit WebSocket
+        try:
+            return await self._client.connect(url, headers=ws_headers, timeout=timeout)
+        except Exception as e:
+            logger.error(f"LivekitWebSocketReverse: Connect failed, {e}")
+            raise UpstreamException(
+                f"LivekitWebSocketReverse: Connect failed, {str(e)}"
+            )
+
+
+__all__ = [
+    "LivekitTokenReverse",
+    "LivekitWebSocketReverse",
+    "LIVEKIT_TOKEN_API",
+    "LIVEKIT_WS_URL",
+]