diff --git a/app/services/grok/protocols/grpc_web.py b/app/services/grok/protocols/grpc_web.py index 0724727..e69de29 100644 --- a/app/services/grok/protocols/grpc_web.py +++ b/app/services/grok/protocols/grpc_web.py @@ -1,157 +0,0 @@ -""" -gRPC-Web 协议工具 - -提供 framing 编码/解码、trailer 解析等通用功能。 -支持 application/grpc-web+proto 和 application/grpc-web-text (base64) 两种格式。 -""" - -from __future__ import annotations - -import base64 -import re -import struct -from dataclasses import dataclass -from typing import Dict, List, Mapping, Tuple -from urllib.parse import unquote - - -_B64_RE = re.compile(rb"^[A-Za-z0-9+/=\r\n]+$") - - -def encode_grpc_web_payload(data: bytes) -> bytes: - """ - 编码 gRPC-Web data frame - - Frame format: - 1-byte flags + 4-byte big-endian length + message bytes - """ - return b"\x00" + struct.pack(">I", len(data)) + data - - -def _maybe_decode_grpc_web_text(body: bytes, content_type: str | None) -> bytes: - """处理 grpc-web-text 模式的 base64 解码""" - ct = (content_type or "").lower() - if "grpc-web-text" in ct: - compact = b"".join(body.split()) - return base64.b64decode(compact, validate=False) - - # 启发式：body 仅包含 base64 字符才尝试解码 - head = body[: min(len(body), 2048)] - if head and _B64_RE.fullmatch(head): - compact = b"".join(body.split()) - try: - return base64.b64decode(compact, validate=True) - except Exception: - return body - return body - - -def _parse_trailer_block(payload: bytes) -> Dict[str, str]: - """解析 trailer frame 内容""" - text = payload.decode("utf-8", errors="replace") - lines = [ln for ln in re.split(r"\r\n|\n", text) if ln] - - trailers: Dict[str, str] = {} - for ln in lines: - if ":" not in ln: - continue - k, v = ln.split(":", 1) - trailers[k.strip().lower()] = v.strip() - - # grpc-message 可能是 percent-encoding - if "grpc-message" in trailers: - trailers["grpc-message"] = unquote(trailers["grpc-message"]) - - return trailers - - -def parse_grpc_web_response( - body: bytes, - content_type: str | None = None, - headers: Mapping[str, str] | None = None, -) -> Tuple[List[bytes], Dict[str, str]]: - """ - 解析 gRPC-Web 响应 - - Returns: - (messages, trailers): data frames 列表和合并后的 trailers - """ - decoded = _maybe_decode_grpc_web_text(body, content_type) - - messages: List[bytes] = [] - trailers: Dict[str, str] = {} - - i = 0 - n = len(decoded) - while i < n: - if n - i < 5: - break - - flag = decoded[i] - length = int.from_bytes(decoded[i + 1 : i + 5], "big") - i += 5 - - if n - i < length: - break - - payload = decoded[i : i + length] - i += length - - if flag & 0x80: # trailer frame - trailers.update(_parse_trailer_block(payload)) - elif flag & 0x01: # compressed (不支持) - raise ValueError("grpc-web compressed flag not supported") - else: - messages.append(payload) - - # 兼容：grpc-status 可能在 response headers 中 - if headers: - lower = {k.lower(): v for k, v in headers.items()} - if "grpc-status" in lower and "grpc-status" not in trailers: - trailers["grpc-status"] = str(lower["grpc-status"]).strip() - if "grpc-message" in lower and "grpc-message" not in trailers: - trailers["grpc-message"] = unquote(str(lower["grpc-message"]).strip()) - - return messages, trailers - - -@dataclass(frozen=True) -class GrpcStatus: - code: int - message: str = "" - - @property - def ok(self) -> bool: - return self.code == 0 - - @property - def http_equiv(self) -> int: - """映射到类 HTTP 状态码""" - mapping = { - 0: 200, # OK - 16: 401, # UNAUTHENTICATED - 7: 403, # PERMISSION_DENIED - 8: 429, # RESOURCE_EXHAUSTED - 4: 504, # DEADLINE_EXCEEDED - 14: 503, # UNAVAILABLE - } - return mapping.get(self.code, 502) - - -def get_grpc_status(trailers: Mapping[str, str]) -> GrpcStatus: - """从 trailers 提取 gRPC 状态""" - raw = str(trailers.get("grpc-status", "")).strip() - msg = str(trailers.get("grpc-message", "")).strip() - try: - code = int(raw) - except Exception: - code = -1 - return GrpcStatus(code=code, message=msg) - - -__all__ = [ - "encode_grpc_web_payload", - "parse_grpc_web_response", - "get_grpc_status", - "GrpcStatus", -] diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 218c334..7313976 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -2,278 +2,10 @@ Grok Imagine WebSocket image service. """ -import asyncio -import certifi -import json -import re -import ssl -import time -import uuid -from typing import AsyncGenerator, Dict, Optional - -import aiohttp - -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" +from app.services.reverse.ws_imagine import ImagineWebSocketReverse -class _BlockedError(Exception): - pass - - -class ImageService: - """Grok Imagine WebSocket image service.""" - - def __init__(self): - self._ssl_context = ssl.create_default_context() - self._ssl_context.load_verify_locations(certifi.where()) - self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)") - - def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]: - proxy_url = get_config("network.base_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) - user_agent = get_config("security.user_agent") - return { - "Cookie": cookie, - "Origin": "https://grok.com", - "User-Agent": user_agent, - "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8", - "Cache-Control": "no-cache", - "Pragma": "no-cache", - } - - def _extract_image_id(self, url: str) -> Optional[str]: - match = self._url_pattern.search(url or "") - return match.group(1) if match else None - - def _is_final_image(self, url: str, blob_size: int) -> bool: - return (url or "").lower().endswith( - (".jpg", ".jpeg") - ) and blob_size > get_config("image.image_ws_final_min_bytes") - - def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]: - if not url or not blob: - return None - - image_id = self._extract_image_id(url) or uuid.uuid4().hex - blob_size = len(blob) - is_final = self._is_final_image(url, blob_size) - - stage = ( - "final" - if is_final - else ( - "medium" - if blob_size > get_config("image.image_ws_medium_min_bytes") - else "preview" - ) - ) - - return { - "type": "image", - "image_id": image_id, - "stage": stage, - "blob": blob, - "blob_size": blob_size, - "url": url, - "is_final": is_final, - } - - async def stream( - self, - token: str, - prompt: str, - aspect_ratio: str = "2:3", - n: int = 1, - enable_nsfw: bool = True, - max_retries: int = None, - ) -> AsyncGenerator[Dict[str, object], None]: - retries = max(1, max_retries if max_retries is not None else 1) - logger.info( - f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}" - ) - - for attempt in range(retries): - try: - yielded_any = False - async for item in self._stream_once( - token, prompt, aspect_ratio, n, enable_nsfw - ): - yielded_any = True - yield item - return - except _BlockedError: - if yielded_any or attempt + 1 >= retries: - if not yielded_any: - yield { - "type": "error", - "error_code": "blocked", - "error": "blocked_no_final_image", - } - return - logger.warning(f"WebSocket blocked, retry {attempt + 1}/{retries}") - except Exception as e: - logger.error(f"WebSocket stream failed: {e}") - return - - async def _stream_once( - self, - token: str, - prompt: str, - aspect_ratio: str, - n: int, - enable_nsfw: bool, - ) -> AsyncGenerator[Dict[str, object], None]: - request_id = str(uuid.uuid4()) - headers = self._get_ws_headers(token) - timeout = float(get_config("network.timeout")) - blocked_seconds = float(get_config("image.image_ws_blocked_seconds")) - - try: - connector, proxy = self._resolve_proxy() - except Exception as e: - logger.error(f"WebSocket proxy setup failed: {e}") - return - - try: - async with aiohttp.ClientSession(connector=connector) as session: - async with session.ws_connect( - WS_URL, - headers=headers, - heartbeat=20, - receive_timeout=timeout, - proxy=proxy, - ) as ws: - message = { - "type": "conversation.item.create", - "timestamp": int(time.time() * 1000), - "item": { - "type": "message", - "content": [ - { - "requestId": request_id, - "text": prompt, - "type": "input_text", - "properties": { - "section_count": 0, - "is_kids_mode": False, - "enable_nsfw": enable_nsfw, - "skip_upsampler": False, - "is_initial": False, - "aspect_ratio": aspect_ratio, - }, - } - ], - }, - } - - await ws.send_json(message) - logger.info(f"WebSocket request sent: {prompt[:80]}...") - - images = {} - completed = 0 - start_time = last_activity = time.time() - medium_received_time = None - - while time.time() - start_time < timeout: - try: - ws_msg = await asyncio.wait_for(ws.receive(), timeout=5.0) - except asyncio.TimeoutError: - if ( - medium_received_time - and completed == 0 - and time.time() - medium_received_time - > min(10, blocked_seconds) - ): - raise _BlockedError() - if completed > 0 and time.time() - last_activity > 10: - logger.info( - f"WebSocket idle timeout, collected {completed} images" - ) - break - continue - - if ws_msg.type == aiohttp.WSMsgType.TEXT: - last_activity = time.time() - msg = json.loads(ws_msg.data) - msg_type = msg.get("type") - - if msg_type == "image": - info = self._classify_image( - msg.get("url", ""), msg.get("blob", "") - ) - if not info: - continue - - image_id = info["image_id"] - existing = images.get(image_id, {}) - - if ( - info["stage"] == "medium" - and medium_received_time is None - ): - medium_received_time = time.time() - - if info["is_final"] and not existing.get("is_final"): - completed += 1 - logger.debug( - f"Final image received: id={image_id}, size={info['blob_size']}" - ) - - images[image_id] = { - "is_final": info["is_final"] - or existing.get("is_final") - } - yield info - - elif msg_type == "error": - logger.warning( - f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}" - ) - yield { - "type": "error", - "error_code": msg.get("err_code", ""), - "error": msg.get("err_msg", ""), - } - return - - if completed >= n: - logger.info( - f"WebSocket collected {completed} final images" - ) - break - - if ( - medium_received_time - and completed == 0 - and time.time() - medium_received_time > blocked_seconds - ): - raise _BlockedError() - - elif ws_msg.type in ( - aiohttp.WSMsgType.CLOSED, - aiohttp.WSMsgType.ERROR, - ): - logger.warning(f"WebSocket closed/error: {ws_msg.type}") - yield { - "type": "error", - "error_code": "ws_closed", - "error": f"websocket closed: {ws_msg.type}", - } - break - - except aiohttp.ClientError as e: - logger.error(f"WebSocket connection error: {e}") - yield {"type": "error", "error_code": "connection_failed", "error": str(e)} - - -image_service = ImageService() +ImageService = ImagineWebSocketReverse +image_service = ImagineWebSocketReverse() __all__ = ["image_service", "ImageService"] diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 208bf95..81515dc 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -22,10 +22,11 @@ class VoiceService: ) -> Dict[str, Any]: browser = get_config("security.browser") async with AsyncSession(impersonate=browser) as session: - return await LivekitTokenReverse.request( + response = await LivekitTokenReverse.request( session, token=token, voice=voice, personality=personality, speed=speed, ) + return response.json() diff --git a/app/services/reverse/__init__.py b/app/services/reverse/__init__.py index 69594b5..08734a8 100644 --- a/app/services/reverse/__init__.py +++ b/app/services/reverse/__init__.py @@ -10,6 +10,7 @@ from .nsfw_mgmt import NsfwMgmtReverse from .rate_limits import RateLimitsReverse from .set_birth import SetBirthReverse from .ws_livekit import LivekitTokenReverse, LivekitWebSocketReverse +from .ws_imagine import ImagineWebSocketReverse from .utils.headers import build_headers from .utils.statsig import StatsigGenerator @@ -25,6 +26,7 @@ __all__ = [ "SetBirthReverse", "LivekitTokenReverse", "LivekitWebSocketReverse", + "ImagineWebSocketReverse", "StatsigGenerator", "build_headers", ] diff --git a/app/services/reverse/utils/grpc.py b/app/services/reverse/utils/grpc.py index 2cb2688..07ee84d 100644 --- a/app/services/reverse/utils/grpc.py +++ b/app/services/reverse/utils/grpc.py @@ -6,7 +6,7 @@ import base64 import re import struct from dataclasses import dataclass -from typing import Dict, List, Mapping, Tuple +from typing import Dict, List, Mapping, Optional, Tuple from urllib.parse import unquote @@ -41,7 +41,7 @@ class GrpcClient: return b"\x00" + struct.pack(">I", len(data)) + data @staticmethod - def _maybe_decode_grpc_web_text(body: bytes, content_type: str | None) -> bytes: + def _maybe_decode_grpc_web_text(body: bytes, content_type: Optional[str]) -> bytes: ct = (content_type or "").lower() if "grpc-web-text" in ct: compact = b"".join(body.split()) @@ -77,8 +77,8 @@ class GrpcClient: def parse_response( cls, body: bytes, - content_type: str | None = None, - headers: Mapping[str, str] | None = None, + content_type: Optional[str] = None, + headers: Optional[Mapping[str, str]] = None, ) -> Tuple[List[bytes], Dict[str, str]]: decoded = cls._maybe_decode_grpc_web_text(body, content_type) @@ -128,26 +128,7 @@ class GrpcClient: return GrpcStatus(code=code, message=msg) -def encode_grpc_web_payload(data: bytes) -> bytes: - return GrpcClient.encode_payload(data) - - -def parse_grpc_web_response( - body: bytes, - content_type: str | None = None, - headers: Mapping[str, str] | None = None, -) -> Tuple[List[bytes], Dict[str, str]]: - return GrpcClient.parse_response(body, content_type=content_type, headers=headers) - - -def get_grpc_status(trailers: Mapping[str, str]) -> GrpcStatus: - return GrpcClient.get_status(trailers) - - __all__ = [ - "encode_grpc_web_payload", - "parse_grpc_web_response", - "get_grpc_status", "GrpcStatus", "GrpcClient", ] diff --git a/app/services/reverse/utils/headers.py b/app/services/reverse/utils/headers.py index 7b41577..03a8f25 100644 --- a/app/services/reverse/utils/headers.py +++ b/app/services/reverse/utils/headers.py @@ -10,9 +10,15 @@ from app.core.config import get_config from app.services.reverse.utils.statsig import StatsigGenerator -def _build_sso_cookie(sso_token: str) -> str: +def build_sso_cookie(sso_token: str) -> str: """ Build SSO Cookie string. + + Args: + sso_token: str, the SSO token. + + Returns: + str: The SSO Cookie string. """ # Format sso_token = sso_token[4:] if sso_token.startswith("sso=") else sso_token @@ -28,20 +34,44 @@ def _build_sso_cookie(sso_token: str) -> str: return cookie -def build_headers( - cookie_token: str, - content_type: Optional[str] = None, - origin: Optional[str] = None, - referer: Optional[str] = None, -) -> Dict[str, str]: +def build_ws_headers(token: Optional[str] = None, origin: Optional[str] = None, extra: Optional[Dict[str, str]] = None) -> Dict[str, str]: + """ + Build headers for WebSocket requests. + + Args: + token: Optional[str], the SSO token for Cookie. Defaults to None. + origin: Optional[str], the Origin value. Defaults to "https://grok.com" if not provided. + extra: Optional[Dict[str, str]], extra headers to merge. Defaults to None. + + Returns: + Dict[str, str]: The headers dictionary. + """ + headers = { + "Origin": origin or "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", + } + + if token: + headers["Cookie"] = build_sso_cookie(token) + + if extra: + headers.update(extra) + + return headers + + +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. + cookie_token: str, the SSO token. + content_type: Optional[str], the Content-Type value. + origin: Optional[str], the Origin value. Defaults to "https://grok.com" if not provided. + referer: Optional[str], the Referer value. Defaults to "https://grok.com/" if not provided. Returns: Dict[str, str]: The headers dictionary. @@ -64,7 +94,7 @@ def build_headers( } # Cookie - headers["Cookie"] = _build_sso_cookie(cookie_token) + headers["Cookie"] = build_sso_cookie(cookie_token) # Content-Type and Accept/Sec-Fetch-Dest if content_type and content_type == "application/json": @@ -101,4 +131,4 @@ def build_headers( return headers -__all__ = ["build_headers"] +__all__ = ["build_headers", "build_sso_cookie", "build_ws_headers"] diff --git a/app/services/reverse/utils/websocket.py b/app/services/reverse/utils/websocket.py index a4a8ffb..67f1516 100644 --- a/app/services/reverse/utils/websocket.py +++ b/app/services/reverse/utils/websocket.py @@ -37,14 +37,12 @@ def _normalize_socks_proxy(proxy_url: str) -> tuple[str, Optional[bool]]: return proxy_url, rdns -def resolve_proxy( - proxy_url: str | None, ssl_context: ssl.SSLContext -) -> tuple[aiohttp.BaseConnector, Optional[str]]: +def resolve_proxy(proxy_url: Optional[str] = None, ssl_context: ssl.SSLContext = _default_ssl_context()) -> tuple[aiohttp.BaseConnector, Optional[str]]: """Resolve proxy connector. Args: - proxy_url: str, the proxy URL. - ssl_context: ssl.SSLContext, the SSL context. + proxy_url: Optional[str], the proxy URL. Defaults to None. + ssl_context: ssl.SSLContext, the SSL context. Defaults to _default_ssl_context(). Returns: tuple[aiohttp.BaseConnector, Optional[str]]: The proxy connector and the proxy URL. @@ -73,11 +71,7 @@ def resolve_proxy( class WebSocketConnection: """WebSocket connection wrapper.""" - def __init__( - self, - session: aiohttp.ClientSession, - ws: aiohttp.ClientWebSocketResponse, - ) -> None: + def __init__(self, session: aiohttp.ClientSession, ws: aiohttp.ClientWebSocketResponse) -> None: self.session = session self.ws = ws @@ -96,22 +90,22 @@ class WebSocketConnection: class WebSocketClient: """WebSocket client with proxy support.""" - def __init__(self, proxy: str | None = None) -> None: + def __init__(self, proxy: Optional[str] = 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, + headers: Optional[Mapping[str, str]] = None, + ws_kwargs: Optional[Mapping[str, object]] = 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. + headers: Optional[Mapping[str, str]], the headers to send. Defaults to None. + ws_kwargs: Optional[Mapping[str, object]], extra ws_connect kwargs. Defaults to None. Returns: WebSocketConnection: The WebSocket connection. @@ -120,20 +114,18 @@ class WebSocketClient: 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) - ) + client_timeout = aiohttp.ClientTimeout(total=get_config("network.timeout")) # Create session session = aiohttp.ClientSession(connector=connector, timeout=client_timeout) try: + extra_kwargs = dict(ws_kwargs or {}) ws = await session.ws_connect( url, headers=headers, proxy=proxy, ssl=self._ssl_context, + **extra_kwargs, ) return WebSocketConnection(session, ws) except Exception: diff --git a/app/services/reverse/ws_imagine.py b/app/services/reverse/ws_imagine.py new file mode 100644 index 0000000..d4c3c7c --- /dev/null +++ b/app/services/reverse/ws_imagine.py @@ -0,0 +1,262 @@ +""" +Reverse interface: Imagine WebSocket image stream. +""" + +import asyncio +import orjson +import re +import time +import uuid +from typing import AsyncGenerator, Dict, Optional + +import aiohttp + +from app.core.config import get_config +from app.core.logger import logger +from app.services.reverse.utils.headers import build_ws_headers +from app.services.reverse.utils.websocket import WebSocketClient + +WS_IMAGINE_URL = "wss://grok.com/ws/imagine/listen" + + +class _BlockedError(Exception): + pass + + +class ImagineWebSocketReverse: + """Imagine WebSocket reverse interface.""" + + def __init__(self) -> None: + self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)") + self._client = WebSocketClient() + + def _extract_image_id(self, url: str) -> Optional[str]: + match = self._url_pattern.search(url or "") + return match.group(1) if match else None + + def _is_final_image(self, url: str, blob_size: int, final_min_bytes: int) -> bool: + return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > final_min_bytes + + def _classify_image(self, url: str, blob: str, final_min_bytes: int, medium_min_bytes: int) -> Optional[Dict[str, object]]: + if not url or not blob: + return None + + image_id = self._extract_image_id(url) or uuid.uuid4().hex + blob_size = len(blob) + is_final = self._is_final_image(url, blob_size, final_min_bytes) + + stage = ( + "final" + if is_final + else ("medium" if blob_size > medium_min_bytes else "preview") + ) + + return { + "type": "image", + "image_id": image_id, + "stage": stage, + "blob": blob, + "blob_size": blob_size, + "url": url, + "is_final": is_final, + } + + def _build_request_message(self, request_id: str, prompt: str, aspect_ratio: str, enable_nsfw: bool) -> Dict[str, object]: + return { + "type": "conversation.item.create", + "timestamp": int(time.time() * 1000), + "item": { + "type": "message", + "content": [ + { + "requestId": request_id, + "text": prompt, + "type": "input_text", + "properties": { + "section_count": 0, + "is_kids_mode": False, + "enable_nsfw": enable_nsfw, + "skip_upsampler": False, + "is_initial": False, + "aspect_ratio": aspect_ratio, + }, + } + ], + }, + } + + async def stream( + self, + token: str, + prompt: str, + aspect_ratio: str = "2:3", + n: int = 1, + enable_nsfw: bool = True, + max_retries: Optional[int] = None, + ) -> AsyncGenerator[Dict[str, object], None]: + retries = max(1, max_retries if max_retries is not None else 1) + logger.info( + f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}" + ) + + for attempt in range(retries): + try: + yielded_any = False + async for item in self._stream_once( + token, prompt, aspect_ratio, n, enable_nsfw + ): + yielded_any = True + yield item + return + except _BlockedError: + if yielded_any or attempt + 1 >= retries: + if not yielded_any: + yield { + "type": "error", + "error_code": "blocked", + "error": "blocked_no_final_image", + } + return + logger.warning(f"WebSocket blocked, retry {attempt + 1}/{retries}") + except Exception as e: + logger.error(f"WebSocket stream failed: {e}") + return + + async def _stream_once( + self, + token: str, + prompt: str, + aspect_ratio: str, + n: int, + enable_nsfw: bool, + ) -> AsyncGenerator[Dict[str, object], None]: + request_id = str(uuid.uuid4()) + headers = build_ws_headers(token=token) + timeout = float(get_config("network.timeout")) + blocked_seconds = float(get_config("image.image_ws_blocked_seconds")) + blocked_grace = min(10.0, blocked_seconds) + final_min_bytes = int(get_config("image.image_ws_final_min_bytes")) + medium_min_bytes = int(get_config("image.image_ws_medium_min_bytes")) + + try: + conn = await self._client.connect( + WS_IMAGINE_URL, + headers=headers, + ws_kwargs={ + "heartbeat": 20, + "receive_timeout": timeout, + }, + ) + except Exception as e: + logger.error(f"WebSocket connect failed: {e}") + yield { + "type": "error", + "error_code": "connection_failed", + "error": str(e), + } + return + + try: + async with conn as ws: + message = self._build_request_message( + request_id, prompt, aspect_ratio, enable_nsfw + ) + await ws.send_json(message) + logger.info(f"WebSocket request sent: {prompt[:80]}...") + + final_ids: set[str] = set() + completed = 0 + start_time = last_activity = time.monotonic() + medium_received_time: Optional[float] = None + + while time.monotonic() - start_time < timeout: + try: + ws_msg = await asyncio.wait_for(ws.receive(), timeout=5.0) + except asyncio.TimeoutError: + now = time.monotonic() + if ( + medium_received_time + and completed == 0 + and now - medium_received_time > blocked_grace + ): + raise _BlockedError() + if completed > 0 and now - last_activity > 10: + logger.info( + f"WebSocket idle timeout, collected {completed} images" + ) + break + continue + + if ws_msg.type == aiohttp.WSMsgType.TEXT: + last_activity = time.monotonic() + try: + msg = orjson.loads(ws_msg.data) + except orjson.JSONDecodeError as e: + logger.warning(f"WebSocket message decode failed: {e}") + continue + + msg_type = msg.get("type") + + if msg_type == "image": + info = self._classify_image( + msg.get("url", ""), + msg.get("blob", ""), + final_min_bytes, + medium_min_bytes, + ) + if not info: + continue + + image_id = info["image_id"] + if info["stage"] == "medium" and medium_received_time is None: + medium_received_time = time.monotonic() + + if info["is_final"] and image_id not in final_ids: + final_ids.add(image_id) + completed += 1 + logger.debug( + f"Final image received: id={image_id}, size={info['blob_size']}" + ) + + yield info + + elif msg_type == "error": + logger.warning( + f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}" + ) + yield { + "type": "error", + "error_code": msg.get("err_code", ""), + "error": msg.get("err_msg", ""), + } + return + + if completed >= n: + logger.info(f"WebSocket collected {completed} final images") + break + + if ( + medium_received_time + and completed == 0 + and time.monotonic() - medium_received_time > blocked_seconds + ): + raise _BlockedError() + + elif ws_msg.type in ( + aiohttp.WSMsgType.CLOSED, + aiohttp.WSMsgType.ERROR, + ): + logger.warning(f"WebSocket closed/error: {ws_msg.type}") + yield { + "type": "error", + "error_code": "ws_closed", + "error": f"websocket closed: {ws_msg.type}", + } + break + + except aiohttp.ClientError as e: + logger.error(f"WebSocket connection error: {e}") + yield {"type": "error", "error_code": "connection_failed", "error": str(e)} + + +__all__ = ["ImagineWebSocketReverse", "WS_IMAGINE_URL"] diff --git a/app/services/reverse/ws_livekit.py b/app/services/reverse/ws_livekit.py index e810fa3..095e31b 100644 --- a/app/services/reverse/ws_livekit.py +++ b/app/services/reverse/ws_livekit.py @@ -11,7 +11,7 @@ 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.headers import build_headers, build_ws_headers from app.services.reverse.utils.retry import retry_on_status from app.services.reverse.utils.websocket import WebSocketClient, WebSocketConnection @@ -29,7 +29,6 @@ class LivekitTokenReverse: voice: str = "ara", personality: str = "assistant", speed: float = 1.0, - livekit_url: str = LIVEKIT_WS_URL, ) -> Dict[str, Any]: """Fetch LiveKit token. @@ -39,7 +38,6 @@ class LivekitTokenReverse: 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. @@ -69,7 +67,7 @@ class LivekitTokenReverse: } ).decode(), "requestAgentDispatch": False, - "livekitUrl": livekit_url, + "livekitUrl": LIVEKIT_WS_URL, "params": {"enable_markdown_transcript": "true"}, } @@ -100,8 +98,7 @@ class LivekitTokenReverse: return response response = await retry_on_status(_do_request) - return response.json() - + return response except Exception as e: # Handle upstream exception @@ -137,108 +134,37 @@ class LivekitWebSocketReverse: 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. + async def connect(self, token: str) -> 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. + token: str, the SSO token. Returns: - str: The LiveKit WebSocket URL. + WebSocketConnection: The LiveKit WebSocket connection. """ - # Build base URL - base = livekit_url.rstrip("/") + # Format URL + base = LIVEKIT_WS_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), + "access_token": token, + "auto_subscribe": "1", + "sdk": "js", + "version": "2.11.4", + "protocol": "15", } - 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, - ) + url = f"{base}?{urlencode(params)}" # Build WebSocket headers - ws_headers = self._build_headers(headers) + ws_headers = build_ws_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) + return await self._client.connect(url, headers=ws_headers) except Exception as e: logger.error(f"LivekitWebSocketReverse: Connect failed, {e}") raise UpstreamException(