feat: enhance configuration management and service initialization

app/api/v1/image.py

@@ -19,7 +19,7 @@ from app.services.grok.services.image import image_service
 from app.services.grok.services.assets import UploadService
 from app.services.grok.services.media import VideoService
 from app.services.grok.models.model import ModelService
-from app.services.grok.processors.processor import (
+from app.services.grok.processors import (
     ImageStreamProcessor,
     ImageCollectProcessor,
     ImageWSStreamProcessor,

app/core/config.py

@@ -53,12 +53,19 @@ class Config:
     def __init__(self):
         self._config = {}
         self._defaults = {}
+        self._code_defaults = {}
         self._defaults_loaded = False
 
+    def register_defaults(self, defaults: Dict[str, Any]):
+        """注册代码中定义的默认值"""
+        self._code_defaults = _deep_merge(self._code_defaults, defaults)
+
     def _ensure_defaults(self):
         if self._defaults_loaded:
             return
-        self._defaults = _load_defaults()
+        file_defaults = _load_defaults()
+        # 合并文件默认值和代码默认值（代码默认值优先级更低）
+        self._defaults = _deep_merge(self._code_defaults, file_defaults)
         self._defaults_loaded = True
 
     async def load(self):
@@ -140,4 +147,9 @@ def get_config(key: str, default: Any = None) -> Any:
     return config.get(key, default)
 
 
-__all__ = ["Config", "config", "get_config"]
+def register_defaults(defaults: Dict[str, Any]):
+    """注册默认配置"""
+    config.register_defaults(defaults)
+
+
+__all__ = ["Config", "config", "get_config", "register_defaults"]

app/services/grok/defaults.py

@@ -0,0 +1,71 @@
+"""
+Grok 服务默认配置
+
+此文件定义所有 Grok 相关服务的默认值，会在应用启动时注册到配置系统中。
+"""
+
+# Grok 服务默认配置
+GROK_DEFAULTS = {
+    "grok": {
+        # 网络配置
+        "browser": "chrome136",
+        "timeout": 120,
+        "user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36",
+        "base_proxy_url": "",
+        "asset_proxy_url": "",
+        "cf_clearance": "",
+        
+        # 对话配置
+        "temporary": True,
+        "disable_memory": True,
+        "stream": True,
+        "thinking": False,
+        "filter_tags": ["grok:render", "xaiartifact", "xai:tool_usage_card"],
+        
+        # 重试配置
+        "max_retry": 3,
+        "retry_status_codes": [401, 429, 403],
+        "retry_backoff_base": 0.5,
+        "retry_backoff_factor": 2.0,
+        "retry_backoff_max": 30.0,
+        "retry_budget": 90.0,
+        
+        # 超时配置
+        "stream_idle_timeout": 45.0,
+        "video_idle_timeout": 90.0,
+        
+        # 图片配置
+        "image_ws_blocked_seconds": 15,
+        "image_ws_final_min_bytes": 100000,
+        "image_ws_medium_min_bytes": 30000,
+        
+        # Statsig
+        "dynamic_statsig": True,
+    },
+    
+    "app": {
+        "app_url": "",
+        "image_format": "url",
+        "video_format": "html",
+    },
+    
+    "performance": {
+        "assets_max_concurrent": 25,
+        "assets_delete_batch_size": 10,
+        "media_max_concurrent": 50,
+        "usage_max_concurrent": 25,
+    },
+    
+    "cache": {
+        "enable_auto_clean": True,
+        "limit_mb": 1024,
+    },
+}
+
+
+def get_grok_defaults():
+    """获取 Grok 默认配置"""
+    return GROK_DEFAULTS
+
+
+__all__ = ["GROK_DEFAULTS", "get_grok_defaults"]

app/services/grok/processors/__init__.py

@@ -0,0 +1,22 @@
+"""
+OpenAI 响应格式处理器
+"""
+
+from .base import BaseProcessor, StreamIdleTimeoutError
+from .chat_processors import StreamProcessor, CollectProcessor
+from .video_processors import VideoStreamProcessor, VideoCollectProcessor
+from .image_processors import ImageStreamProcessor, ImageCollectProcessor
+from .image_ws_processors import ImageWSStreamProcessor, ImageWSCollectProcessor
+
+__all__ = [
+    "BaseProcessor",
+    "StreamIdleTimeoutError",
+    "StreamProcessor",
+    "CollectProcessor",
+    "VideoStreamProcessor",
+    "VideoCollectProcessor",
+    "ImageStreamProcessor",
+    "ImageCollectProcessor",
+    "ImageWSStreamProcessor",
+    "ImageWSCollectProcessor",
+]

app/services/grok/processors/base.py

@@ -0,0 +1,159 @@
+"""
+响应处理器基类和通用工具
+"""
+
+import asyncio
+import time
+from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar
+
+from app.core.config import get_config
+from app.core.logger import logger
+from app.services.grok.services.assets import DownloadService
+
+
+ASSET_URL = "https://assets.grok.com/"
+
+T = TypeVar("T")
+
+
+def _is_http2_stream_error(e: Exception) -> bool:
+    """检查是否为 HTTP/2 流错误"""
+    err_str = str(e).lower()
+    return "http/2" in err_str or "curl: (92)" in err_str or "stream" in err_str
+
+
+def _normalize_stream_line(line: Any) -> Optional[str]:
+    """规范化流式响应行，兼容 SSE data 前缀与空行"""
+    if line is None:
+        return None
+    if isinstance(line, (bytes, bytearray)):
+        text = line.decode("utf-8", errors="ignore")
+    else:
+        text = str(line)
+    text = text.strip()
+    if not text:
+        return None
+    if text.startswith("data:"):
+        text = text[5:].strip()
+    if text == "[DONE]":
+        return None
+    return text
+
+
+def _collect_image_urls(obj: Any) -> List[str]:
+    """递归收集响应中的图片 URL"""
+    urls: List[str] = []
+    seen = set()
+
+    def add(url: str):
+        if not url or url in seen:
+            return
+        seen.add(url)
+        urls.append(url)
+
+    def walk(value: Any):
+        if isinstance(value, dict):
+            for key, item in value.items():
+                if key in {"generatedImageUrls", "imageUrls", "imageURLs"}:
+                    if isinstance(item, list):
+                        for url in item:
+                            if isinstance(url, str):
+                                add(url)
+                    elif isinstance(item, str):
+                        add(item)
+                    continue
+                walk(item)
+        elif isinstance(value, list):
+            for item in value:
+                walk(item)
+
+    walk(obj)
+    return urls
+
+
+class StreamIdleTimeoutError(Exception):
+    """流空闲超时错误"""
+
+    def __init__(self, idle_seconds: float):
+        self.idle_seconds = idle_seconds
+        super().__init__(f"Stream idle timeout after {idle_seconds}s")
+
+
+async def _with_idle_timeout(
+    iterable: AsyncIterable[T], idle_timeout: float, model: str = ""
+) -> AsyncGenerator[T, None]:
+    """
+    包装异步迭代器，添加空闲超时检测
+
+    Args:
+        iterable: 原始异步迭代器
+        idle_timeout: 空闲超时时间(秒)，0 表示禁用
+        model: 模型名称(用于日志)
+    """
+    if idle_timeout <= 0:
+        async for item in iterable:
+            yield item
+        return
+
+    iterator = iterable.__aiter__()
+    while True:
+        try:
+            item = await asyncio.wait_for(iterator.__anext__(), timeout=idle_timeout)
+            yield item
+        except asyncio.TimeoutError:
+            logger.warning(
+                f"Stream idle timeout after {idle_timeout}s",
+                extra={"model": model, "idle_timeout": idle_timeout},
+            )
+            raise StreamIdleTimeoutError(idle_timeout)
+        except StopAsyncIteration:
+            break
+
+
+class BaseProcessor:
+    """基础处理器"""
+
+    def __init__(self, model: str, token: str = ""):
+        self.model = model
+        self.token = token
+        self.created = int(time.time())
+        self.app_url = get_config("app.app_url")
+        self._dl_service: Optional[DownloadService] = None
+
+    def _get_dl(self) -> DownloadService:
+        """获取下载服务实例（复用）"""
+        if self._dl_service is None:
+            self._dl_service = DownloadService()
+        return self._dl_service
+
+    async def close(self):
+        """释放下载服务资源"""
+        if self._dl_service:
+            await self._dl_service.close()
+            self._dl_service = None
+
+    async def process_url(self, path: str, media_type: str = "image") -> str:
+        """处理资产 URL"""
+        if path.startswith("http"):
+            from urllib.parse import urlparse
+            path = urlparse(path).path
+
+        if not path.startswith("/"):
+            path = f"/{path}"
+
+        if self.app_url:
+            dl_service = self._get_dl()
+            await dl_service.download(path, self.token, media_type)
+            return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}"
+        else:
+            return f"{ASSET_URL.rstrip('/')}{path}"
+
+
+__all__ = [
+    "BaseProcessor",
+    "StreamIdleTimeoutError",
+    "_with_idle_timeout",
+    "_normalize_stream_line",
+    "_collect_image_urls",
+    "_is_http2_stream_error",
+]

app/services/grok/processors/chat_processors.py

@@ -0,0 +1,326 @@
+"""
+聊天响应处理器
+"""
+
+import asyncio
+import uuid
+import re
+from typing import Any, AsyncGenerator, AsyncIterable
+
+import orjson
+from curl_cffi.requests.errors import RequestsError
+
+from app.core.config import get_config
+from app.core.logger import logger
+from app.core.exceptions import UpstreamException
+from .base import (
+    BaseProcessor,
+    StreamIdleTimeoutError,
+    _with_idle_timeout,
+    _normalize_stream_line,
+    _collect_image_urls,
+    _is_http2_stream_error,
+)
+
+
+class StreamProcessor(BaseProcessor):
+    """流式响应处理器"""
+
+    def __init__(self, model: str, token: str = "", think: bool = None):
+        super().__init__(model, token)
+        self.response_id: str = None
+        self.fingerprint: str = ""
+        self.think_opened: bool = False
+        self.role_sent: bool = False
+        self.filter_tags = get_config("grok.filter_tags")
+        self.image_format = get_config("app.image_format")
+        self._tag_buffer: str = ""
+        self._in_filter_tag: bool = False
+
+        if think is None:
+            self.show_think = get_config("grok.thinking")
+        else:
+            self.show_think = think
+
+    def _filter_token(self, token: str) -> str:
+        """过滤 token 中的特殊标签（如 <grok:render>...</grok:render>），支持跨 token 的标签过滤"""
+        if not self.filter_tags:
+            return token
+
+        result = []
+        i = 0
+        while i < len(token):
+            char = token[i]
+
+            if self._in_filter_tag:
+                self._tag_buffer += char
+                if char == ">":
+                    if "/>" in self._tag_buffer:
+                        self._in_filter_tag = False
+                        self._tag_buffer = ""
+                    else:
+                        for tag in self.filter_tags:
+                            if f"</{tag}>" in self._tag_buffer:
+                                self._in_filter_tag = False
+                                self._tag_buffer = ""
+                                break
+                i += 1
+                continue
+
+            if char == "<":
+                remaining = token[i:]
+                tag_started = False
+                for tag in self.filter_tags:
+                    if remaining.startswith(f"<{tag}"):
+                        tag_started = True
+                        break
+                    if len(remaining) < len(tag) + 1:
+                        for j in range(1, len(remaining) + 1):
+                            if f"<{tag}".startswith(remaining[:j]):
+                                tag_started = True
+                                break
+
+                if tag_started:
+                    self._in_filter_tag = True
+                    self._tag_buffer = char
+                    i += 1
+                    continue
+
+            result.append(char)
+            i += 1
+
+        return "".join(result)
+
+    def _sse(self, content: str = "", role: str = None, finish: str = None) -> str:
+        """构建 SSE 响应"""
+        delta = {}
+        if role:
+            delta["role"] = role
+            delta["content"] = ""
+        elif content:
+            delta["content"] = content
+
+        chunk = {
+            "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}",
+            "object": "chat.completion.chunk",
+            "created": self.created,
+            "model": self.model,
+            "system_fingerprint": self.fingerprint,
+            "choices": [
+                {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}
+            ],
+        }
+        return f"data: {orjson.dumps(chunk).decode()}\n\n"
+
+    async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
+        """处理流式响应"""
+        idle_timeout = get_config("grok.stream_idle_timeout")
+
+        try:
+            async for line in _with_idle_timeout(response, idle_timeout, self.model):
+                line = _normalize_stream_line(line)
+                if not line:
+                    continue
+                try:
+                    data = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    continue
+
+                resp = data.get("result", {}).get("response", {})
+
+                if (llm := resp.get("llmInfo")) and not self.fingerprint:
+                    self.fingerprint = llm.get("modelHash", "")
+                if rid := resp.get("responseId"):
+                    self.response_id = rid
+
+                if not self.role_sent:
+                    yield self._sse(role="assistant")
+                    self.role_sent = True
+
+                # 图像生成进度
+                if img := resp.get("streamingImageGenerationResponse"):
+                    if self.show_think:
+                        if not self.think_opened:
+                            yield self._sse("<think>\n")
+                            self.think_opened = True
+                        idx = img.get("imageIndex", 0) + 1
+                        progress = img.get("progress", 0)
+                        yield self._sse(f"正在生成第{idx}张图片中，当前进度{progress}%\n")
+                    continue
+
+                # modelResponse
+                if mr := resp.get("modelResponse"):
+                    if self.think_opened and self.show_think:
+                        if msg := mr.get("message"):
+                            yield self._sse(msg + "\n")
+                        yield self._sse("</think>\n")
+                        self.think_opened = False
+
+                    # 处理生成的图片
+                    for url in _collect_image_urls(mr):
+                        parts = url.split("/")
+                        img_id = parts[-2] if len(parts) >= 2 else "image"
+
+                        if self.image_format == "base64":
+                            dl_service = self._get_dl()
+                            base64_data = await dl_service.to_base64(url, self.token, "image")
+                            if base64_data:
+                                yield self._sse(f"![{img_id}]({base64_data})\n")
+                            else:
+                                final_url = await self.process_url(url, "image")
+                                yield self._sse(f"![{img_id}]({final_url})\n")
+                        else:
+                            final_url = await self.process_url(url, "image")
+                            yield self._sse(f"![{img_id}]({final_url})\n")
+
+                    if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"):
+                        self.fingerprint = meta["llm_info"]["modelHash"]
+                    continue
+
+                # 普通 token
+                if (token := resp.get("token")) is not None:
+                    if token:
+                        filtered = self._filter_token(token)
+                        if filtered:
+                            yield self._sse(filtered)
+
+            if self.think_opened:
+                yield self._sse("</think>\n")
+            yield self._sse(finish="stop")
+            yield "data: [DONE]\n\n"
+        except asyncio.CancelledError:
+            logger.debug("Stream cancelled by client", extra={"model": self.model})
+        except StreamIdleTimeoutError as e:
+            raise UpstreamException(
+                message=f"Stream idle timeout after {e.idle_seconds}s",
+                status_code=504,
+                details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds},
+            )
+        except RequestsError as e:
+            if _is_http2_stream_error(e):
+                logger.warning(f"HTTP/2 stream error: {e}", extra={"model": self.model})
+                raise UpstreamException(
+                    message="Upstream connection closed unexpectedly",
+                    status_code=502,
+                    details={"error": str(e), "type": "http2_stream_error"},
+                )
+            logger.error(f"Stream request error: {e}", extra={"model": self.model})
+            raise UpstreamException(
+                message=f"Upstream request failed: {e}",
+                status_code=502,
+                details={"error": str(e)},
+            )
+        except Exception as e:
+            logger.error(f"Stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
+            raise
+        finally:
+            await self.close()
+
+
+class CollectProcessor(BaseProcessor):
+    """非流式响应处理器"""
+
+    def __init__(self, model: str, token: str = ""):
+        super().__init__(model, token)
+        self.image_format = get_config("app.image_format")
+        self.filter_tags = get_config("grok.filter_tags")
+
+    def _filter_content(self, content: str) -> str:
+        """过滤内容中的特殊标签"""
+        if not content or not self.filter_tags:
+            return content
+
+        result = content
+        for tag in self.filter_tags:
+            pattern = rf"<{re.escape(tag)}[^>]*>.*?</{re.escape(tag)}>|<{re.escape(tag)}[^>]*/>"
+            result = re.sub(pattern, "", result, flags=re.DOTALL)
+
+        return result
+
+    async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]:
+        """处理并收集完整响应"""
+        response_id = ""
+        fingerprint = ""
+        content = ""
+        idle_timeout = get_config("grok.stream_idle_timeout")
+
+        try:
+            async for line in _with_idle_timeout(response, idle_timeout, self.model):
+                line = _normalize_stream_line(line)
+                if not line:
+                    continue
+                try:
+                    data = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    continue
+
+                resp = data.get("result", {}).get("response", {})
+
+                if (llm := resp.get("llmInfo")) and not fingerprint:
+                    fingerprint = llm.get("modelHash", "")
+
+                if mr := resp.get("modelResponse"):
+                    response_id = mr.get("responseId", "")
+                    content = mr.get("message", "")
+
+                    if urls := _collect_image_urls(mr):
+                        content += "\n"
+                        for url in urls:
+                            parts = url.split("/")
+                            img_id = parts[-2] if len(parts) >= 2 else "image"
+
+                            if self.image_format == "base64":
+                                dl_service = self._get_dl()
+                                base64_data = await dl_service.to_base64(url, self.token, "image")
+                                if base64_data:
+                                    content += f"![{img_id}]({base64_data})\n"
+                                else:
+                                    final_url = await self.process_url(url, "image")
+                                    content += f"![{img_id}]({final_url})\n"
+                            else:
+                                final_url = await self.process_url(url, "image")
+                                content += f"![{img_id}]({final_url})\n"
+
+                    if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"):
+                        fingerprint = meta["llm_info"]["modelHash"]
+
+        except asyncio.CancelledError:
+            logger.debug("Collect cancelled by client", extra={"model": self.model})
+        except StreamIdleTimeoutError as e:
+            logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model})
+        except RequestsError as e:
+            if _is_http2_stream_error(e):
+                logger.warning(f"HTTP/2 stream error in collect: {e}", extra={"model": self.model})
+            else:
+                logger.error(f"Collect request error: {e}", extra={"model": self.model})
+        except Exception as e:
+            logger.error(f"Collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
+        finally:
+            await self.close()
+
+        content = self._filter_content(content)
+
+        return {
+            "id": response_id,
+            "object": "chat.completion",
+            "created": self.created,
+            "model": self.model,
+            "system_fingerprint": fingerprint,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content, "refusal": None, "annotations": []},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0,
+                "prompt_tokens_details": {"cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0},
+                "completion_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0},
+            },
+        }
+
+
+__all__ = ["StreamProcessor", "CollectProcessor"]

app/services/grok/processors/image_processors.py

@@ -0,0 +1,211 @@
+"""
+图片生成响应处理器（HTTP）
+"""
+
+import asyncio
+import random
+from typing import Any, AsyncGenerator, AsyncIterable, List
+
+import orjson
+from curl_cffi.requests.errors import RequestsError
+
+from app.core.config import get_config
+from app.core.logger import logger
+from app.core.exceptions import UpstreamException
+from .base import (
+    BaseProcessor,
+    StreamIdleTimeoutError,
+    _with_idle_timeout,
+    _normalize_stream_line,
+    _collect_image_urls,
+    _is_http2_stream_error,
+)
+
+
+class ImageStreamProcessor(BaseProcessor):
+    """图片生成流式响应处理器"""
+
+    def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json"):
+        super().__init__(model, token)
+        self.partial_index = 0
+        self.n = n
+        self.target_index = random.randint(0, 1) if n == 1 else None
+        self.response_format = response_format
+        if response_format == "url":
+            self.response_field = "url"
+        elif response_format == "base64":
+            self.response_field = "base64"
+        else:
+            self.response_field = "b64_json"
+
+    def _sse(self, event: str, data: dict) -> str:
+        """构建 SSE 响应"""
+        return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n"
+
+    async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
+        """处理流式响应"""
+        final_images = []
+        idle_timeout = get_config("grok.stream_idle_timeout")
+
+        try:
+            async for line in _with_idle_timeout(response, idle_timeout, self.model):
+                line = _normalize_stream_line(line)
+                if not line:
+                    continue
+                try:
+                    data = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    continue
+
+                resp = data.get("result", {}).get("response", {})
+
+                # 图片生成进度
+                if img := resp.get("streamingImageGenerationResponse"):
+                    image_index = img.get("imageIndex", 0)
+                    progress = img.get("progress", 0)
+
+                    if self.n == 1 and image_index != self.target_index:
+                        continue
+
+                    out_index = 0 if self.n == 1 else image_index
+
+                    yield self._sse(
+                        "image_generation.partial_image",
+                        {
+                            "type": "image_generation.partial_image",
+                            self.response_field: "",
+                            "index": out_index,
+                            "progress": progress,
+                        },
+                    )
+                    continue
+
+                # modelResponse
+                if mr := resp.get("modelResponse"):
+                    if urls := _collect_image_urls(mr):
+                        for url in urls:
+                            if self.response_format == "url":
+                                processed = await self.process_url(url, "image")
+                                if processed:
+                                    final_images.append(processed)
+                                continue
+                            dl_service = self._get_dl()
+                            base64_data = await dl_service.to_base64(url, self.token, "image")
+                            if base64_data:
+                                if "," in base64_data:
+                                    b64 = base64_data.split(",", 1)[1]
+                                else:
+                                    b64 = base64_data
+                                final_images.append(b64)
+                    continue
+
+            for index, b64 in enumerate(final_images):
+                if self.n == 1:
+                    if index != self.target_index:
+                        continue
+                    out_index = 0
+                else:
+                    out_index = index
+
+                yield self._sse(
+                    "image_generation.completed",
+                    {
+                        "type": "image_generation.completed",
+                        self.response_field: b64,
+                        "index": out_index,
+                        "usage": {
+                            "total_tokens": 0,
+                            "input_tokens": 0,
+                            "output_tokens": 0,
+                            "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
+                        },
+                    },
+                )
+        except asyncio.CancelledError:
+            logger.debug("Image stream cancelled by client")
+        except StreamIdleTimeoutError as e:
+            raise UpstreamException(
+                message=f"Image stream idle timeout after {e.idle_seconds}s",
+                status_code=504,
+                details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds},
+            )
+        except RequestsError as e:
+            if _is_http2_stream_error(e):
+                logger.warning(f"HTTP/2 stream error in image: {e}")
+                raise UpstreamException(
+                    message="Upstream connection closed unexpectedly",
+                    status_code=502,
+                    details={"error": str(e), "type": "http2_stream_error"},
+                )
+            logger.error(f"Image stream request error: {e}")
+            raise UpstreamException(
+                message=f"Upstream request failed: {e}",
+                status_code=502,
+                details={"error": str(e)},
+            )
+        except Exception as e:
+            logger.error(f"Image stream processing error: {e}", extra={"error_type": type(e).__name__})
+            raise
+        finally:
+            await self.close()
+
+
+class ImageCollectProcessor(BaseProcessor):
+    """图片生成非流式响应处理器"""
+
+    def __init__(self, model: str, token: str = "", response_format: str = "b64_json"):
+        super().__init__(model, token)
+        self.response_format = response_format
+
+    async def process(self, response: AsyncIterable[bytes]) -> List[str]:
+        """处理并收集图片"""
+        images = []
+        idle_timeout = get_config("grok.stream_idle_timeout")
+
+        try:
+            async for line in _with_idle_timeout(response, idle_timeout, self.model):
+                line = _normalize_stream_line(line)
+                if not line:
+                    continue
+                try:
+                    data = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    continue
+
+                resp = data.get("result", {}).get("response", {})
+
+                if mr := resp.get("modelResponse"):
+                    if urls := _collect_image_urls(mr):
+                        for url in urls:
+                            if self.response_format == "url":
+                                processed = await self.process_url(url, "image")
+                                if processed:
+                                    images.append(processed)
+                                continue
+                            dl_service = self._get_dl()
+                            base64_data = await dl_service.to_base64(url, self.token, "image")
+                            if base64_data:
+                                if "," in base64_data:
+                                    b64 = base64_data.split(",", 1)[1]
+                                else:
+                                    b64 = base64_data
+                                images.append(b64)
+
+        except asyncio.CancelledError:
+            logger.debug("Image collect cancelled by client")
+        except StreamIdleTimeoutError as e:
+            logger.warning(f"Image collect idle timeout: {e}")
+        except RequestsError as e:
+            if _is_http2_stream_error(e):
+                logger.warning(f"HTTP/2 stream error in image collect: {e}")
+            else:
+                logger.error(f"Image collect request error: {e}")
+        except Exception as e:
+            logger.error(f"Image collect processing error: {e}", extra={"error_type": type(e).__name__})
+        finally:
+            await self.close()
+
+        return images
+
+
+__all__ = ["ImageStreamProcessor", "ImageCollectProcessor"]

app/services/grok/processors/image_ws_processors.py

@@ -0,0 +1,231 @@
+"""
+图片生成响应处理器（WebSocket）
+"""
+
+import base64
+import time
+from pathlib import Path
+from typing import AsyncGenerator, AsyncIterable, List, Dict, Optional
+
+import orjson
+
+from app.core.config import get_config
+from app.core.logger import logger
+from app.core.exceptions import UpstreamException
+from .base import BaseProcessor
+
+
+class ImageWSBaseProcessor(BaseProcessor):
+    """WebSocket 图片处理基类"""
+
+    def __init__(self, model: str, token: str = "", response_format: str = "b64_json"):
+        super().__init__(model, token)
+        self.response_format = response_format
+        if response_format == "url":
+            self.response_field = "url"
+        elif response_format == "base64":
+            self.response_field = "base64"
+        else:
+            self.response_field = "b64_json"
+        self._image_dir: Optional[Path] = None
+
+    def _ensure_image_dir(self) -> Path:
+        if self._image_dir is None:
+            base_dir = Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" / "image"
+            base_dir.mkdir(parents=True, exist_ok=True)
+            self._image_dir = base_dir
+        return self._image_dir
+
+    def _strip_base64(self, blob: str) -> str:
+        if not blob:
+            return ""
+        if "," in blob and "base64" in blob.split(",", 1)[0]:
+            return blob.split(",", 1)[1]
+        return blob
+
+    def _filename(self, image_id: str, is_final: bool) -> str:
+        ext = "jpg" if is_final else "png"
+        return f"{image_id}.{ext}"
+
+    def _build_file_url(self, filename: str) -> str:
+        app_url = get_config("app.app_url")
+        if app_url:
+            return f"{app_url.rstrip('/')}/v1/files/image/{filename}"
+        return f"/v1/files/image/{filename}"
+
+    def _save_blob(self, image_id: str, blob: str, is_final: bool) -> str:
+        data = self._strip_base64(blob)
+        if not data:
+            return ""
+        image_dir = self._ensure_image_dir()
+        filename = self._filename(image_id, is_final)
+        filepath = image_dir / filename
+        with open(filepath, "wb") as f:
+            f.write(base64.b64decode(data))
+        return self._build_file_url(filename)
+
+    def _pick_best(self, existing: Optional[Dict], incoming: Dict) -> Dict:
+        if not existing:
+            return incoming
+        if incoming.get("is_final") and not existing.get("is_final"):
+            return incoming
+        if existing.get("is_final") and not incoming.get("is_final"):
+            return existing
+        if incoming.get("blob_size", 0) > existing.get("blob_size", 0):
+            return incoming
+        return existing
+
+    def _to_output(self, image_id: str, item: Dict) -> str:
+        try:
+            if self.response_format == "url":
+                return self._save_blob(image_id, item.get("blob", ""), item.get("is_final", False))
+            return self._strip_base64(item.get("blob", ""))
+        except Exception as e:
+            logger.warning(f"Image output failed: {e}")
+            return ""
+
+
+class ImageWSStreamProcessor(ImageWSBaseProcessor):
+    """WebSocket 图片流式响应处理器"""
+
+    def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json", size: str = "1024x1024"):
+        super().__init__(model, token, "b64_json")
+        self.n = n
+        self.size = size
+        self._target_id: Optional[str] = None
+        self._index_map: Dict[str, int] = {}
+        self._partial_map: Dict[str, int] = {}
+
+    def _assign_index(self, image_id: str) -> Optional[int]:
+        if image_id in self._index_map:
+            return self._index_map[image_id]
+        if len(self._index_map) >= self.n:
+            return None
+        self._index_map[image_id] = len(self._index_map)
+        return self._index_map[image_id]
+
+    def _sse(self, event: str, data: dict) -> str:
+        return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n"
+
+    async def process(self, response: AsyncIterable[dict]) -> AsyncGenerator[str, None]:
+        images: Dict[str, Dict] = {}
+
+        async for item in response:
+            if item.get("type") == "error":
+                message = item.get("error") or "Upstream error"
+                code = item.get("error_code") or "upstream_error"
+                yield self._sse("error", {"error": {"message": message, "type": "server_error", "code": code}})
+                return
+            if item.get("type") != "image":
+                continue
+
+            image_id = item.get("image_id")
+            if not image_id:
+                continue
+
+            if self.n == 1:
+                if self._target_id is None:
+                    self._target_id = image_id
+                index = 0 if image_id == self._target_id else None
+            else:
+                index = self._assign_index(image_id)
+
+            images[image_id] = self._pick_best(images.get(image_id), item)
+
+            if index is None:
+                continue
+
+            if item.get("stage") != "final":
+                partial_b64 = self._strip_base64(item.get("blob", ""))
+                if not partial_b64:
+                    continue
+                partial_index = self._partial_map.get(image_id, 0)
+                if item.get("stage") == "medium":
+                    partial_index = max(partial_index, 1)
+                self._partial_map[image_id] = partial_index
+                yield self._sse(
+                    "image_generation.partial_image",
+                    {
+                        "type": "image_generation.partial_image",
+                        "b64_json": partial_b64,
+                        "created_at": int(time.time()),
+                        "size": self.size,
+                        "index": index,
+                        "partial_image_index": partial_index,
+                    },
+                )
+
+        if self.n == 1:
+            if self._target_id and self._target_id in images:
+                selected = [(self._target_id, images[self._target_id])]
+            else:
+                selected = [max(images.items(), key=lambda x: (x[1].get("is_final", False), x[1].get("blob_size", 0)))] if images else []
+        else:
+            selected = [(image_id, images[image_id]) for image_id in self._index_map if image_id in images]
+
+        for image_id, item in selected:
+            output = self._strip_base64(item.get("blob", ""))
+            if not output:
+                continue
+
+            if self.n == 1:
+                index = 0
+            else:
+                index = self._index_map.get(image_id, 0)
+            yield self._sse(
+                "image_generation.completed",
+                {
+                    "type": "image_generation.completed",
+                    "b64_json": output,
+                    "created_at": int(time.time()),
+                    "size": self.size,
+                    "index": index,
+                    "usage": {
+                        "total_tokens": 0,
+                        "input_tokens": 0,
+                        "output_tokens": 0,
+                        "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
+                    },
+                },
+            )
+
+
+class ImageWSCollectProcessor(ImageWSBaseProcessor):
+    """WebSocket 图片非流式响应处理器"""
+
+    def __init__(self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json"):
+        super().__init__(model, token, response_format)
+        self.n = n
+
+    async def process(self, response: AsyncIterable[dict]) -> List[str]:
+        images: Dict[str, Dict] = {}
+
+        async for item in response:
+            if item.get("type") == "error":
+                message = item.get("error") or "Upstream error"
+                raise UpstreamException(message, details=item)
+            if item.get("type") != "image":
+                continue
+            image_id = item.get("image_id")
+            if not image_id:
+                continue
+            images[image_id] = self._pick_best(images.get(image_id), item)
+
+        selected = sorted(
+            images.values(),
+            key=lambda x: (x.get("is_final", False), x.get("blob_size", 0)),
+            reverse=True,
+        )
+        if self.n:
+            selected = selected[: self.n]
+
+        results: List[str] = []
+        for item in selected:
+            output = self._to_output(item.get("image_id", ""), item)
+            if output:
+                results.append(output)
+
+        return results
+
+
+__all__ = ["ImageWSStreamProcessor", "ImageWSCollectProcessor"]

app/services/grok/processors/video_processors.py

@@ -0,0 +1,236 @@
+"""
+视频响应处理器
+"""
+
+import asyncio
+import uuid
+from typing import Any, AsyncGenerator, AsyncIterable, Optional
+
+import orjson
+from curl_cffi.requests.errors import RequestsError
+
+from app.core.config import get_config
+from app.core.logger import logger
+from app.core.exceptions import UpstreamException
+from .base import (
+    BaseProcessor,
+    StreamIdleTimeoutError,
+    _with_idle_timeout,
+    _normalize_stream_line,
+    _is_http2_stream_error,
+)
+
+
+class VideoStreamProcessor(BaseProcessor):
+    """视频流式响应处理器"""
+
+    def __init__(self, model: str, token: str = "", think: bool = None):
+        super().__init__(model, token)
+        self.response_id: Optional[str] = None
+        self.think_opened: bool = False
+        self.role_sent: bool = False
+        self.video_format = str(get_config("app.video_format")).lower()
+
+        if think is None:
+            self.show_think = get_config("grok.thinking")
+        else:
+            self.show_think = think
+
+    def _sse(self, content: str = "", role: str = None, finish: str = None) -> str:
+        """构建 SSE 响应"""
+        delta = {}
+        if role:
+            delta["role"] = role
+            delta["content"] = ""
+        elif content:
+            delta["content"] = content
+
+        chunk = {
+            "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}",
+            "object": "chat.completion.chunk",
+            "created": self.created,
+            "model": self.model,
+            "choices": [
+                {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}
+            ],
+        }
+        return f"data: {orjson.dumps(chunk).decode()}\n\n"
+
+    def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str:
+        """构建视频 HTML 标签"""
+        import html
+        safe_video_url = html.escape(video_url)
+        safe_thumbnail_url = html.escape(thumbnail_url)
+        poster_attr = f' poster="{safe_thumbnail_url}"' if safe_thumbnail_url else ""
+        return f'''<video id="video" controls="" preload="none"{poster_attr}>
+  <source id="mp4" src="{safe_video_url}" type="video/mp4">
+</video>'''
+
+    async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
+        """处理视频流式响应"""
+        idle_timeout = get_config("grok.video_idle_timeout")
+
+        try:
+            async for line in _with_idle_timeout(response, idle_timeout, self.model):
+                line = _normalize_stream_line(line)
+                if not line:
+                    continue
+                try:
+                    data = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    continue
+
+                resp = data.get("result", {}).get("response", {})
+
+                if rid := resp.get("responseId"):
+                    self.response_id = rid
+
+                if not self.role_sent:
+                    yield self._sse(role="assistant")
+                    self.role_sent = True
+
+                # 视频生成进度
+                if video_resp := resp.get("streamingVideoGenerationResponse"):
+                    progress = video_resp.get("progress", 0)
+
+                    if self.show_think:
+                        if not self.think_opened:
+                            yield self._sse("<think>\n")
+                            self.think_opened = True
+                        yield self._sse(f"正在生成视频中，当前进度{progress}%\n")
+
+                    if progress == 100:
+                        video_url = video_resp.get("videoUrl", "")
+                        thumbnail_url = video_resp.get("thumbnailImageUrl", "")
+
+                        if self.think_opened and self.show_think:
+                            yield self._sse("</think>\n")
+                            self.think_opened = False
+
+                        if video_url:
+                            final_video_url = await self.process_url(video_url, "video")
+                            final_thumbnail_url = ""
+                            if thumbnail_url:
+                                final_thumbnail_url = await self.process_url(thumbnail_url, "image")
+
+                            if self.video_format == "url":
+                                yield self._sse(final_video_url)
+                            else:
+                                video_html = self._build_video_html(final_video_url, final_thumbnail_url)
+                                yield self._sse(video_html)
+
+                            logger.info(f"Video generated: {video_url}")
+                    continue
+
+            if self.think_opened:
+                yield self._sse("</think>\n")
+            yield self._sse(finish="stop")
+            yield "data: [DONE]\n\n"
+        except asyncio.CancelledError:
+            logger.debug("Video stream cancelled by client", extra={"model": self.model})
+        except StreamIdleTimeoutError as e:
+            raise UpstreamException(
+                message=f"Video stream idle timeout after {e.idle_seconds}s",
+                status_code=504,
+                details={"error": str(e), "type": "stream_idle_timeout", "idle_seconds": e.idle_seconds},
+            )
+        except RequestsError as e:
+            if _is_http2_stream_error(e):
+                logger.warning(f"HTTP/2 stream error in video: {e}", extra={"model": self.model})
+                raise UpstreamException(
+                    message="Upstream connection closed unexpectedly",
+                    status_code=502,
+                    details={"error": str(e), "type": "http2_stream_error"},
+                )
+            logger.error(f"Video stream request error: {e}", extra={"model": self.model})
+            raise UpstreamException(
+                message=f"Upstream request failed: {e}",
+                status_code=502,
+                details={"error": str(e)},
+            )
+        except Exception as e:
+            logger.error(f"Video stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
+        finally:
+            await self.close()
+
+
+class VideoCollectProcessor(BaseProcessor):
+    """视频非流式响应处理器"""
+
+    def __init__(self, model: str, token: str = ""):
+        super().__init__(model, token)
+        self.video_format = str(get_config("app.video_format")).lower()
+
+    def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str:
+        poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else ""
+        return f'''<video id="video" controls="" preload="none"{poster_attr}>
+  <source id="mp4" src="{video_url}" type="video/mp4">
+</video>'''
+
+    async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]:
+        """处理并收集视频响应"""
+        response_id = ""
+        content = ""
+        idle_timeout = get_config("grok.video_idle_timeout")
+
+        try:
+            async for line in _with_idle_timeout(response, idle_timeout, self.model):
+                line = _normalize_stream_line(line)
+                if not line:
+                    continue
+                try:
+                    data = orjson.loads(line)
+                except orjson.JSONDecodeError:
+                    continue
+
+                resp = data.get("result", {}).get("response", {})
+
+                if video_resp := resp.get("streamingVideoGenerationResponse"):
+                    if video_resp.get("progress") == 100:
+                        response_id = resp.get("responseId", "")
+                        video_url = video_resp.get("videoUrl", "")
+                        thumbnail_url = video_resp.get("thumbnailImageUrl", "")
+
+                        if video_url:
+                            final_video_url = await self.process_url(video_url, "video")
+                            final_thumbnail_url = ""
+                            if thumbnail_url:
+                                final_thumbnail_url = await self.process_url(thumbnail_url, "image")
+
+                            if self.video_format == "url":
+                                content = final_video_url
+                            else:
+                                content = self._build_video_html(final_video_url, final_thumbnail_url)
+                            logger.info(f"Video generated: {video_url}")
+
+        except asyncio.CancelledError:
+            logger.debug("Video collect cancelled by client", extra={"model": self.model})
+        except StreamIdleTimeoutError as e:
+            logger.warning(f"Video collect idle timeout: {e}", extra={"model": self.model})
+        except RequestsError as e:
+            if _is_http2_stream_error(e):
+                logger.warning(f"HTTP/2 stream error in video collect: {e}", extra={"model": self.model})
+            else:
+                logger.error(f"Video collect request error: {e}", extra={"model": self.model})
+        except Exception as e:
+            logger.error(f"Video collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__})
+        finally:
+            await self.close()
+
+        return {
+            "id": response_id,
+            "object": "chat.completion",
+            "created": self.created,
+            "model": self.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content, "refusal": None},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+        }
+
+
+__all__ = ["VideoStreamProcessor", "VideoCollectProcessor"]

app/services/grok/services/chat.py

@@ -10,39 +10,22 @@ from curl_cffi.requests import AsyncSession
 
 from app.core.logger import logger
 from app.core.config import get_config
-from app.core.exceptions import (
-    AppException,
-    UpstreamException,
-    ValidationException,
-    ErrorType,
-)
+from app.core.exceptions import AppException, UpstreamException, ValidationException, ErrorType
 from app.services.grok.models.model import ModelService
 from app.services.grok.services.assets import UploadService
-from app.services.grok.processors.processor import StreamProcessor, CollectProcessor
+from app.services.grok.processors import StreamProcessor, CollectProcessor
 from app.services.grok.utils.retry import retry_on_status
 from app.services.grok.utils.headers import apply_statsig, build_sso_cookie
+from app.services.grok.utils.stream import wrap_stream_with_usage
 from app.services.token import get_token_manager, EffortType
 
 
 CHAT_API = "https://grok.com/rest/app-chat/conversations/new"
 
-DEFAULT_TIMEOUT = 120
-DEFAULT_BROWSER = "chrome136"
-DEFAULT_USER_AGENT = (
-    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
-    "AppleWebKit/537.36 (KHTML, like Gecko) "
-    "Chrome/136.0.0.0 Safari/537.36"
-)
-DEFAULT_TEMPORARY = True
-DEFAULT_DISABLE_MEMORY = True
-DEFAULT_STREAM = True
-DEFAULT_BASE_PROXY_URL = ""
-
 
 @dataclass
 class ChatRequest:
     """聊天请求数据"""
-
     model: str
     messages: List[Dict[str, Any]]
     stream: bool = None
@@ -52,85 +35,47 @@ class ChatRequest:
 class MessageExtractor:
     """消息内容提取器"""
 
-    # 需要上传的类型
-    UPLOAD_TYPES = {"image_url", "input_audio", "file"}
-    # 视频模式不支持的类型
-    VIDEO_UNSUPPORTED = {"input_audio", "file"}
-
     @staticmethod
-    def extract(
-        messages: List[Dict[str, Any]], is_video: bool = False
-    ) -> tuple[str, List[tuple[str, str]]]:
-        """
-        从 OpenAI 消息格式提取内容
-
-        Args:
-            messages: OpenAI 格式消息列表
-            is_video: 是否为视频模型
-
-        Returns:
-            (text, attachments): 拼接后的文本和需要上传的附件列表
-
-        Raises:
-            ValueError: 视频模型遇到不支持的内容类型
-        """
+    def extract(messages: List[Dict[str, Any]], is_video: bool = False) -> tuple[str, List[tuple[str, str]]]:
+        """从 OpenAI 消息格式提取内容，返回 (text, attachments)"""
         texts = []
-        attachments = []  # 需要上传的附件 (URL 或 base64)
-
-        # 先抽取每条消息的文本，保留角色信息用于合并
-        extracted: List[Dict[str, str]] = []
+        attachments = []
+        extracted = []
 
         for msg in messages:
             role = msg.get("role", "")
             content = msg.get("content", "")
             parts = []
 
-            # 简单字符串内容
             if isinstance(content, str):
                 if content.strip():
                     parts.append(content)
-
-            # 列表格式内容
             elif isinstance(content, list):
                 for item in content:
                     item_type = item.get("type", "")
 
-                    # 文本类型
                     if item_type == "text":
-                        text = item.get("text", "")
-                        if text.strip():
+                        if text := item.get("text", "").strip():
                             parts.append(text)
 
-                    # 图片类型
                     elif item_type == "image_url":
                         image_data = item.get("image_url", {})
-                        url = (
-                            image_data.get("url", "")
-                            if isinstance(image_data, dict)
-                            else str(image_data)
-                        )
+                        url = image_data.get("url", "") if isinstance(image_data, dict) else str(image_data)
                         if url:
                             attachments.append(("image", url))
 
-                    # 音频类型
                     elif item_type == "input_audio":
                         if is_video:
                             raise ValueError("视频模型不支持 input_audio 类型")
                         audio_data = item.get("input_audio", {})
-                        data = (
-                            audio_data.get("data", "")
-                            if isinstance(audio_data, dict)
-                            else str(audio_data)
-                        )
+                        data = audio_data.get("data", "") if isinstance(audio_data, dict) else str(audio_data)
                         if data:
                             attachments.append(("audio", data))
 
-                    # 文件类型
                     elif item_type == "file":
                         if is_video:
                             raise ValueError("视频模型不支持 file 类型")
                         file_data = item.get("file", {})
-                        # file 可能是 URL 或 base64
                         url = file_data.get("url", "") or file_data.get("data", "")
                         if isinstance(file_data, str):
                             url = file_data
@@ -140,30 +85,15 @@ class MessageExtractor:
             if parts:
                 extracted.append({"role": role, "text": "\n".join(parts)})
 
-        # 合并文本
-        last_user_index = None
-        for i in range(len(extracted) - 1, -1, -1):
-            if extracted[i]["role"] == "user":
-                last_user_index = i
-                break
+        # 找到最后一条 user 消息
+        last_user_index = next((i for i in range(len(extracted) - 1, -1, -1) if extracted[i]["role"] == "user"), None)
 
         for i, item in enumerate(extracted):
             role = item["role"] or "user"
             text = item["text"]
-            if i == last_user_index:
-                texts.append(text)
-            else:
-                texts.append(f"{role}: {text}")
+            texts.append(text if i == last_user_index else f"{role}: {text}")
 
-        # 换行拼接文本
-        message = "\n\n".join(texts)
-        return message, attachments
-
-    @staticmethod
-    def extract_text_only(messages: List[Dict[str, Any]]) -> str:
-        """仅提取文本内容"""
-        text, _ = MessageExtractor.extract(messages, is_video=True)
-        return text
+        return "\n\n".join(texts), attachments
 
 
 class ChatRequestBuilder:
@@ -172,7 +102,7 @@ class ChatRequestBuilder:
     @staticmethod
     def build_headers(token: str) -> Dict[str, str]:
         """构造请求头"""
-        user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT)
+        user_agent = get_config("grok.user_agent")
         headers = {
             "Accept": "*/*",
             "Accept-Encoding": "gzip, deflate, br, zstd",
@@ -202,33 +132,17 @@ class ChatRequestBuilder:
         return headers
 
     @staticmethod
-    def build_payload(
-        message: str,
-        model: str,
-        mode: str = None,
-        file_attachments: List[str] = None,
-        image_attachments: List[str] = None,
-    ) -> Dict[str, Any]:
-        """
-        构造请求体
-
-        Args:
-            message: 消息文本
-            model: 模型名称
-            mode: 模型模式
-            file_attachments: 文件附件 ID 列表
-            image_attachments: 图片附件 ID 列表（合并到 fileAttachments）
-        """
-        temporary = get_config("grok.temporary", DEFAULT_TEMPORARY)
-
-        merged_attachments: List[str] = []
+    def build_payload(message: str, model: str, mode: str = None, 
+                     file_attachments: List[str] = None, image_attachments: List[str] = None) -> Dict[str, Any]:
+        """构造请求体"""
+        merged_attachments = []
         if file_attachments:
             merged_attachments.extend(file_attachments)
         if image_attachments:
             merged_attachments.extend(image_attachments)
 
-        payload: Dict[str, Any] = {
-            "temporary": temporary,
+        payload = {
+            "temporary": get_config("grok.temporary"),
             "modelName": model,
             "message": message,
             "fileAttachments": merged_attachments,
@@ -236,23 +150,17 @@ class ChatRequestBuilder:
             "disableSearch": False,
             "enableImageGeneration": True,
             "returnImageBytes": False,
-            "returnRawGrokInXaiRequest": False,
             "enableImageStreaming": True,
             "imageGenerationCount": 2,
             "forceConcise": False,
             "toolOverrides": {},
             "enableSideBySide": True,
             "sendFinalMetadata": True,
-            "isReasoning": False,
-            "disableTextFollowUps": False,
             "responseMetadata": {
                 "modelConfigOverride": {"modelMap": {}},
                 "requestModelDetails": {"modelId": model},
             },
-            "disableMemory": get_config("grok.disable_memory", DEFAULT_DISABLE_MEMORY),
-            "forceSideBySide": False,
-            "isAsyncChat": False,
-            "disableSelfHarmShortCircuit": False,
+            "disableMemory": get_config("grok.disable_memory"),
             "deviceEnvInfo": {
                 "darkModeEnabled": False,
                 "devicePixelRatio": 2,
@@ -269,70 +177,31 @@ class ChatRequestBuilder:
         return payload
 
 
-# ==================== Grok 服务 ====================
-
-
 class GrokChatService:
     """Grok API 调用服务"""
 
     def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL)
+        self.proxy = proxy or get_config("grok.base_proxy_url")
 
-    async def chat(
-        self,
-        token: str,
-        message: str,
-        model: str = "grok-3",
-        mode: str = None,
-        stream: bool = None,
-        file_attachments: List[str] = None,
-        image_attachments: List[str] = None,
-        raw_payload: Dict[str, Any] = None,
-    ):
-        """
-        发送聊天请求
-
-        Args:
-            token: 认证 Token
-            message: 消息文本
-            model: Grok 模型名称
-            mode: 模型模式
-            stream: 是否流式
-            file_attachments: 文件附件 ID 列表
-            image_attachments: 图片附件 URL 列表
-
-        Raises:
-            UpstreamException: 当 Grok API 返回错误且重试耗尽时
-        """
+    async def chat(self, token: str, message: str, model: str = "grok-3", mode: str = None,
+                   stream: bool = None, file_attachments: List[str] = None, 
+                   image_attachments: List[str] = None, raw_payload: Dict[str, Any] = None):
+        """发送聊天请求"""
         if stream is None:
-            stream = get_config("grok.stream", DEFAULT_STREAM)
+            stream = get_config("grok.stream")
 
         headers = ChatRequestBuilder.build_headers(token)
-        if raw_payload is not None:
-            payload = raw_payload
-        else:
-            payload = ChatRequestBuilder.build_payload(
-                message,
-                model,
-                mode,
-                file_attachments,
-                image_attachments,
-            )
+        payload = raw_payload if raw_payload is not None else ChatRequestBuilder.build_payload(
+            message, model, mode, file_attachments, image_attachments
+        )
         proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None
-        timeout = get_config("grok.timeout", DEFAULT_TIMEOUT)
-
-        # 状态码提取器
-        def extract_status(e: Exception) -> int | None:
-            if isinstance(e, UpstreamException) and e.details:
-                return e.details.get("status")
-            return None
+        timeout = get_config("grok.timeout")
 
         logger.debug(f"Payload: {payload}")
         
-        # 建立连接函数
+        # 建立连接
         async def establish_connection():
-            """建立连接并返回 response 对象"""
-            browser = get_config("grok.browser", DEFAULT_BROWSER)
+            browser = get_config("grok.browser")
             session = AsyncSession(impersonate=browser)
             try:
                 response = await session.post(
@@ -345,63 +214,44 @@ class GrokChatService:
                 )
 
                 if response.status_code != 200:
+                    content = ""
                     try:
                         content = await response.text()
-                    except Exception:
-                        content = ""
-
-                    logger.error(
-                        f"Chat failed: {response.status_code}",
-                        extra={
-                            "status": response.status_code,
-                            "token": token[:10] + "...",
-                        },
-                    )
-                    body_for_log = content if content else "<empty>"
-                    logger.debug(
-                        "Grok API error response "
-                        f"(status={response.status_code}): {body_for_log}"
-                    )
-                    # 关闭 session 并抛出异常
-                    try:
-                        await session.close()
                     except Exception:
                         pass
+
+                    logger.error(f"Chat failed: {response.status_code}", 
+                               extra={"status": response.status_code, "token": token[:10] + "..."})
+                    
+                    await session.close()
                     raise UpstreamException(
                         message=f"Grok API request failed: {response.status_code}",
-                        details={
-                            "status": response.status_code,
-                            "body": content,
-                        },
+                        details={"status": response.status_code, "body": content},
                     )
 
-                # 返回 session 和 response
                 return session, response
 
             except UpstreamException:
-                # 已经处理过的异常，直接抛出
                 raise
             except Exception as e:
-                # 其他异常，关闭 session 并包装
                 logger.error(f"Chat request error: {e}")
-                try:
-                    await session.close()
-                except Exception:
-                    pass
+                await session.close()
                 raise UpstreamException(
                     message=f"Chat connection failed: {str(e)}",
                     details={"error": str(e)},
                 )
 
-        # 建立连接
+        # 重试机制
+        def extract_status(e: Exception) -> int | None:
+            if isinstance(e, UpstreamException) and e.details:
+                return e.details.get("status")
+            return None
+
         session = None
         response = None
         try:
-            session, response = await retry_on_status(
-                establish_connection, extract_status=extract_status
-            )
+            session, response = await retry_on_status(establish_connection, extract_status=extract_status)
         except Exception as e:
-            # 记录失败
             status_code = extract_status(e)
             if status_code:
                 token_mgr = await get_token_manager()
@@ -436,124 +286,47 @@ class GrokChatService:
 
         # 提取消息和附件
         try:
-            message, attachments = MessageExtractor.extract(
-                request.messages, is_video=is_video
-            )
+            message, attachments = MessageExtractor.extract(request.messages, is_video=is_video)
         except ValueError as e:
             raise ValidationException(str(e))
 
-        # 处理附件上传
+        # 上传附件
         file_ids = []
-
         if attachments:
             upload_service = UploadService()
             try:
-                for attach_type, attach_data in attachments:
-                    # 获取 ID
+                for _, attach_data in attachments:
                     file_id, _ = await upload_service.upload(attach_data, token)
-
-                    if attach_type == "image":
-                        file_ids.append(file_id)
-                        logger.debug(f"Image uploaded: {file_id}")
-                    else:
-                        file_ids.append(file_id)
-                        logger.debug(f"File uploaded: {file_id}")
+                    file_ids.append(file_id)
             finally:
                 await upload_service.close()
 
-        stream = (
-            request.stream
-            if request.stream is not None
-            else get_config("grok.stream", DEFAULT_STREAM)
-        )
+        stream = request.stream if request.stream is not None else get_config("grok.stream")
 
         response = await self.chat(
-            token,
-            message,
-            grok_model,
-            mode,
-            stream,
-            file_attachments=file_ids,
-            image_attachments=[],
+            token, message, grok_model, mode, stream,
+            file_attachments=file_ids, image_attachments=[]
         )
 
         return response, stream, request.model
 
 
-# ==================== Chat 业务服务 ====================
-
-
 class ChatService:
     """Chat 业务服务"""
 
     @staticmethod
-    async def _wrap_stream(stream: AsyncGenerator, token_mgr, token: str, model: str):
-        """
-        包装流式响应，在完成时记录使用
-
-        Args:
-            stream: 原始 AsyncGenerator
-            token_mgr: TokenManager 实例
-            token: Token 字符串
-            model: 模型名称
-        """
-        success = False
-        try:
-            async for chunk in stream:
-                yield chunk
-            success = True
-        finally:
-            # 只在成功完成时记录使用，失败/异常时不扣费
-            if success:
-                try:
-                    model_info = ModelService.get(model)
-                    effort = (
-                        EffortType.HIGH
-                        if (model_info and model_info.cost.value == "high")
-                        else EffortType.LOW
-                    )
-                    await token_mgr.consume(token, effort)
-                    logger.debug(
-                        f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})"
-                    )
-                except Exception as e:
-                    logger.warning(f"Failed to record stream usage: {e}")
-
-    @staticmethod
-    async def completions(
-        model: str,
-        messages: List[Dict[str, Any]],
-        stream: bool = None,
-        thinking: str = None,
-    ):
-        """
-        Chat Completions 入口
-
-        Args:
-            model: 模型名称
-            messages: 消息列表
-            stream: 是否流式
-            thinking: 思考模式
-
-        Returns:
-            AsyncGenerator 或 dict
-        """
+    async def completions(model: str, messages: List[Dict[str, Any]], 
+                         stream: bool = None, thinking: str = None):
+        """Chat Completions 入口"""
         # 获取 token
-        try:
-            token_mgr = await get_token_manager()
-            await token_mgr.reload_if_stale()
-            token = None
-            for pool_name in ModelService.pool_candidates_for_model(model):
-                token = token_mgr.get_token(pool_name)
-                if token:
-                    break
-        except Exception as e:
-            logger.error(f"Failed to get token: {e}")
-            raise AppException(
-                message="Internal service error obtaining token",
-                error_type=ErrorType.SERVER.value,
-                code="internal_error",
-            )
+        token_mgr = await get_token_manager()
+        await token_mgr.reload_if_stale()
+        
+        token = None
+        for pool_name in ModelService.pool_candidates_for_model(model):
+            token = token_mgr.get_token(pool_name)
+            if token:
+                break
 
         if not token:
             raise AppException(
@@ -564,61 +337,30 @@ class ChatService:
             )
 
         # 解析参数
-        think = None
-        if thinking == "enabled":
-            think = True
-        elif thinking == "disabled":
-            think = False
-
-        is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM)
+        think = {"enabled": True, "disabled": False}.get(thinking)
+        is_stream = stream if stream is not None else get_config("grok.stream")
 
         # 构造请求
-        chat_request = ChatRequest(
-            model=model, messages=messages, stream=is_stream, think=think
-        )
+        chat_request = ChatRequest(model=model, messages=messages, stream=is_stream, think=think)
 
         # 请求 Grok
         service = GrokChatService()
-        try:
-            response, _, model_name = await service.chat_openai(token, chat_request)
-        except AppException:
-            raise
-        except Exception as e:
-            logger.error(f"Chat service error: {e}")
-            raise UpstreamException(
-                message=f"Service processing failed: {str(e)}",
-                details={"error": str(e)},
-            )
+        response, _, model_name = await service.chat_openai(token, chat_request)
 
         # 处理响应
         if is_stream:
             processor = StreamProcessor(model_name, token, think)
-            return ChatService._wrap_stream(
-                processor.process(response), token_mgr, token, model
-            )
-        else:
-            result = await CollectProcessor(model_name, token).process(response)
-            # 非流式：处理完成后立即记录使用
-            try:
-                model_info = ModelService.get(model)
-                effort = (
-                    EffortType.HIGH
-                    if (model_info and model_info.cost.value == "high")
-                    else EffortType.LOW
-                )
-                await token_mgr.consume(token, effort)
-                logger.debug(
-                    f"Collect completed, recorded usage for token {token[:10]}... (effort={effort.value})"
-                )
-            except Exception as e:
-                logger.warning(f"Failed to record collect usage: {e}")
-            return result
+            return wrap_stream_with_usage(processor.process(response), token_mgr, token, model)
+        
+        # 非流式
+        result = await CollectProcessor(model_name, token).process(response)
+        try:
+            model_info = ModelService.get(model)
+            effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW
+            await token_mgr.consume(token, effort)
+        except Exception as e:
+            logger.warning(f"Failed to record usage: {e}")
+        return result
 
 
-__all__ = [
-    "GrokChatService",
-    "ChatRequest",
-    "ChatRequestBuilder",
-    "MessageExtractor",
-    "ChatService",
-]
+__all__ = ["GrokChatService", "ChatRequest", "ChatRequestBuilder", "MessageExtractor", "ChatService"]

app/services/grok/services/image.py

@@ -18,25 +18,12 @@ from app.core.config import get_config
 from app.core.logger import logger
 from app.services.grok.utils.headers import build_sso_cookie
 
-
 WS_URL = "wss://grok.com/ws/imagine/listen"
 
-DEFAULT_TIMEOUT = 120
-DEFAULT_BLOCKED_SECONDS = 15
-DEFAULT_FINAL_MIN_BYTES = 100000
-DEFAULT_MEDIUM_MIN_BYTES = 30000
-DEFAULT_USER_AGENT = (
-    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
-    "AppleWebKit/537.36 (KHTML, like Gecko) "
-    "Chrome/136.0.0.0 Safari/537.36"
-)
-DEFAULT_BASE_PROXY_URL = ""
-
 
 class _BlockedError(Exception):
     pass
 
-
 class ImageService:
     """Grok Imagine WebSocket image service."""
 
@@ -44,25 +31,22 @@ class ImageService:
         self._ssl_context = ssl.create_default_context()
         self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)")
 
-    def _get_ws_url(self) -> str:
-        return WS_URL
-
     def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]:
-        proxy_url = get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL)
+        proxy_url = get_config("grok.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"Grok Imagine WebSocket using SOCKS proxy: {proxy_url}")
+            logger.info(f"Using SOCKS proxy: {proxy_url}")
             return ProxyConnector.from_url(proxy_url, ssl=self._ssl_context), None
 
-        logger.info(f"Grok Imagine WebSocket using HTTP proxy: {proxy_url}")
+        logger.info(f"Using HTTP proxy: {proxy_url}")
         return aiohttp.TCPConnector(ssl=self._ssl_context), proxy_url
 
     def _get_ws_headers(self, token: str) -> Dict[str, str]:
         cookie = build_sso_cookie(token, include_rw=True)
-        user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT)
+        user_agent = get_config("grok.user_agent")
         return {
             "Cookie": cookie,
             "Origin": "https://grok.com",
@@ -74,15 +58,10 @@ class ImageService:
 
     def _extract_image_id(self, url: str) -> Optional[str]:
         match = self._url_pattern.search(url or "")
-        if match:
-            return match.group(1)
-        return None
+        return match.group(1) if match else None
 
     def _is_final_image(self, url: str, blob_size: int) -> bool:
-        min_bytes = int(
-            get_config("grok.image_ws_final_min_bytes", DEFAULT_FINAL_MIN_BYTES)
-        )
-        return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > min_bytes
+        return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > get_config("grok.image_ws_final_min_bytes")
 
     def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]:
         if not url or not blob:
@@ -91,12 +70,8 @@ class ImageService:
         image_id = self._extract_image_id(url) or uuid.uuid4().hex
         blob_size = len(blob)
         is_final = self._is_final_image(url, blob_size)
-        if is_final:
-            stage = "final"
-        elif blob_size > DEFAULT_MEDIUM_MIN_BYTES:
-            stage = "medium"
-        else:
-            stage = "preview"
+        
+        stage = "final" if is_final else ("medium" if blob_size > get_config("grok.image_ws_medium_min_bytes") else "preview")
 
         return {
             "type": "image",
@@ -117,37 +92,23 @@ class ImageService:
         enable_nsfw: bool = True,
         max_retries: int = None,
     ) -> AsyncGenerator[Dict[str, object], None]:
-        if max_retries is None:
-            retries = 1
-        else:
-            retries = max_retries
-        retries = max(1, retries)
+        retries = max(1, max_retries if max_retries is not None else 1)
 
         for attempt in range(retries):
             try:
                 yielded_any = False
-                async for item in self._stream_once(
-                    token, prompt, aspect_ratio, n, enable_nsfw
-                ):
+                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:
+                if yielded_any or attempt + 1 >= retries:
+                    if not yielded_any:
+                        yield {"type": "error", "error_code": "blocked", "error": "blocked_no_final_image"}
                     return
-                if attempt + 1 < retries:
-                    logger.warning(
-                        f"Grok Imagine WebSocket blocked, retry {attempt + 1}/{retries}"
-                    )
-                    continue
-                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"Grok Imagine WebSocket stream failed: {e}")
+                logger.error(f"WebSocket stream failed: {e}")
                 return
 
     async def _stream_once(
@@ -158,25 +119,21 @@ class ImageService:
         n: int,
         enable_nsfw: bool,
     ) -> AsyncGenerator[Dict[str, object], None]:
-        ws_url = self._get_ws_url()
-
         request_id = str(uuid.uuid4())
         headers = self._get_ws_headers(token)
-        timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT))
-        blocked_seconds = float(
-            get_config("grok.image_ws_blocked_seconds", DEFAULT_BLOCKED_SECONDS)
-        )
+        timeout = float(get_config("grok.timeout"))
+        blocked_seconds = float(get_config("grok.image_ws_blocked_seconds"))
 
         try:
             connector, proxy = self._resolve_proxy()
         except Exception as e:
-            logger.error(f"Grok Imagine WebSocket proxy setup failed: {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,
+                    WS_URL,
                     headers=headers,
                     heartbeat=20,
                     receive_timeout=timeout,
@@ -206,27 +163,21 @@ class ImageService:
                     }
 
                     await ws.send_json(message)
-                    logger.info(f"Grok Imagine WebSocket request sent: {prompt[:80]}...")
+                    logger.info(f"WebSocket request sent: {prompt[:80]}...")
 
-                    images: Dict[str, Dict[str, object]] = {}
+                    images = {}
                     completed = 0
-                    start_time = time.time()
-                    last_activity = time.time()
+                    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:
-                                if time.time() - medium_received_time > min(
-                                    10, blocked_seconds
-                                ):
-                                    raise _BlockedError()
+                            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"Grok Imagine WebSocket idle timeout, collected {completed} images"
-                                )
+                                logger.info(f"WebSocket idle timeout, collected {completed} images")
                                 break
                             continue
 
@@ -236,73 +187,43 @@ class ImageService:
                             msg_type = msg.get("type")
 
                             if msg_type == "image":
-                                info = self._classify_image(
-                                    msg.get("url", ""), msg.get("blob", "")
-                                )
+                                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
 
-                                images[image_id] = {
-                                    "is_final": info["is_final"] or existing.get("is_final")
-                                }
-
+                                images[image_id] = {"is_final": info["is_final"] or existing.get("is_final")}
                                 yield info
 
                             elif msg_type == "error":
-                                error_code = msg.get("err_code", "")
-                                error_msg = msg.get("err_msg", "")
-                                logger.warning(
-                                    f"Grok Imagine WebSocket error: {error_code} - {error_msg}"
-                                )
-                                yield {
-                                    "type": "error",
-                                    "error_code": error_code,
-                                    "error": error_msg,
-                                }
+                                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"Grok Imagine WebSocket collected {completed} final images"
-                                )
+                                logger.info(f"WebSocket collected {completed} final images")
                                 break
 
-                            if medium_received_time and completed == 0:
-                                if time.time() - medium_received_time > blocked_seconds:
-                                    raise _BlockedError()
+                            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"Grok Imagine WebSocket closed/error: {ws_msg.type}"
-                            )
-                            yield {
-                                "type": "error",
-                                "error_code": "ws_closed",
-                                "error": f"websocket closed: {ws_msg.type}",
-                            }
+                        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"Grok Imagine WebSocket connection error: {e}")
-            yield {
-                "type": "error",
-                "error_code": "connection_failed",
-                "error": str(e),
-            }
-
+            logger.error(f"WebSocket connection error: {e}")
+            yield {"type": "error", "error_code": "connection_failed", "error": str(e)}
 
 image_service = ImageService()
 
-
 __all__ = ["image_service", "ImageService"]

app/services/grok/services/media.py

@@ -18,39 +18,39 @@ from app.core.exceptions import (
 )
 from app.services.grok.models.model import ModelService
 from app.services.token import get_token_manager, EffortType
-from app.services.grok.processors.processor import VideoStreamProcessor, VideoCollectProcessor
+from app.services.grok.processors import VideoStreamProcessor, VideoCollectProcessor
 from app.services.grok.utils.headers import apply_statsig, build_sso_cookie
-
+from app.services.grok.utils.stream import wrap_stream_with_usage
 
 CREATE_POST_API = "https://grok.com/rest/media/post/create"
 CHAT_API = "https://grok.com/rest/app-chat/conversations/new"
 
-DEFAULT_BROWSER = "chrome136"
-DEFAULT_USER_AGENT = (
-    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
-    "AppleWebKit/537.36 (KHTML, like Gecko) "
-    "Chrome/136.0.0.0 Safari/537.36"
-)
-DEFAULT_TIMEOUT = 300
-DEFAULT_MAX_CONCURRENT = 50
-DEFAULT_BASE_PROXY_URL = ""
-DEFAULT_STREAM = True
-_MEDIA_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT)
-_MEDIA_SEM_VALUE = DEFAULT_MAX_CONCURRENT
+_MEDIA_SEMAPHORE = None
+_MEDIA_SEM_VALUE = 0
+
+
+def _get_semaphore() -> asyncio.Semaphore:
+    """获取或更新信号量"""
+    global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE
+    value = max(1, int(get_config("performance.media_max_concurrent")))
+    if value != _MEDIA_SEM_VALUE:
+        _MEDIA_SEM_VALUE = value
+        _MEDIA_SEMAPHORE = asyncio.Semaphore(value)
+    return _MEDIA_SEMAPHORE
 
 
 class VideoService:
     """视频生成服务"""
 
     def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL)
-        self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT)
+        self.proxy = proxy or get_config("grok.base_proxy_url")
+        self.timeout = get_config("grok.timeout")
 
     def _build_headers(
         self, token: str, referer: str = "https://grok.com/imagine"
     ) -> dict:
         """构建请求头"""
-        user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT)
+        user_agent = get_config("grok.user_agent")
         headers = {
             "Accept": "*/*",
             "Accept-Encoding": "gzip, deflate, br, zstd",
@@ -84,27 +84,14 @@ class VideoService:
         return {"http": self.proxy, "https": self.proxy} if self.proxy else None
 
     async def create_post(
-        self,
-        token: str,
-        prompt: str,
+        self, token: str, prompt: str,
         media_type: str = "MEDIA_POST_TYPE_VIDEO",
-        media_url: str = None,
+        media_url: str = None
     ) -> str:
-        """
-        创建媒体帖子
-
-        Args:
-            token: 认证 Token
-            prompt: 提示词（视频生成用）
-            media_type: 媒体类型 (MEDIA_POST_TYPE_VIDEO 或 MEDIA_POST_TYPE_IMAGE)
-            media_url: 媒体 URL（图片模式用）
-
-        Returns:
-            post ID
-        """
+        """创建媒体帖子，返回 post ID"""
         try:
             headers = self._build_headers(token)
-
+            
             # 根据类型构建不同的载荷
             if media_type == "MEDIA_POST_TYPE_IMAGE" and media_url:
                 payload = {"mediaType": media_type, "mediaUrl": media_url}
@@ -112,352 +99,160 @@ class VideoService:
                 payload = {"mediaType": media_type, "prompt": prompt}
 
             async with AsyncSession() as session:
-                browser = get_config("grok.browser", DEFAULT_BROWSER)
                 response = await session.post(
                     CREATE_POST_API,
                     headers=headers,
                     json=payload,
-                    impersonate=browser,
+                    impersonate=get_config("grok.browser"),
                     timeout=30,
                     proxies=self._build_proxies(),
                 )
 
             if response.status_code != 200:
                 logger.error(f"Create post failed: {response.status_code}")
-                raise UpstreamException(
-                    f"Failed to create post: {response.status_code}"
-                )
-
-            data = response.json()
-            post_id = data.get("post", {}).get("id", "")
+                raise UpstreamException(f"Failed to create post: {response.status_code}")
 
+            post_id = response.json().get("post", {}).get("id", "")
             if not post_id:
                 raise UpstreamException("No post ID in response")
 
             logger.info(f"Media post created: {post_id} (type={media_type})")
             return post_id
 
+        except AppException:
+            raise
         except Exception as e:
             logger.error(f"Create post error: {e}")
-            if isinstance(e, AppException):
-                raise e
             raise UpstreamException(f"Create post error: {str(e)}")
 
     async def create_image_post(self, token: str, image_url: str) -> str:
-        """
-        创建图片帖子
-
-        Args:
-            token: 认证 Token
-            image_url: 完整的图片 URL (https://assets.grok.com/...)
-
-        Returns:
-            post ID
-        """
-        return await self.create_post(
-            token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url
-        )
+        """创建图片帖子，返回 post ID"""
+        return await self.create_post(token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url)
 
     def _build_payload(
-        self,
-        prompt: str,
-        post_id: str,
+        self, prompt: str, post_id: str,
         aspect_ratio: str = "3:2",
         video_length: int = 6,
         resolution_name: str = "480p",
-        preset: str = "normal",
+        preset: str = "normal"
     ) -> dict:
         """构建视频生成载荷"""
-        mode_flag = "--mode=custom"
-        if preset == "fun":
-            mode_flag = "--mode=extremely-crazy"
-        elif preset == "normal":
-            mode_flag = "--mode=normal"
-        elif preset == "spicy":
-            mode_flag = "--mode=extremely-spicy-or-crazy"
-
-        full_prompt = f"{prompt} {mode_flag}"
-        video_gen_config = {
-            "parentPostId": post_id,
-            "aspectRatio": aspect_ratio,
-            "videoLength": video_length,
-            "resolutionName": resolution_name,
-        }
+        mode_map = {"fun": "--mode=extremely-crazy", "normal": "--mode=normal", "spicy": "--mode=extremely-spicy-or-crazy"}
+        mode_flag = mode_map.get(preset, "--mode=custom")
+        
         return {
             "temporary": True,
             "modelName": "grok-3",
-            "message": full_prompt,
+            "message": f"{prompt} {mode_flag}",
             "toolOverrides": {"videoGen": True},
             "enableSideBySide": True,
             "responseMetadata": {
                 "experiments": [],
                 "modelConfigOverride": {
                     "modelMap": {
-                        "videoGenModelConfig": video_gen_config
+                        "videoGenModelConfig": {
+                            "parentPostId": post_id,
+                            "aspectRatio": aspect_ratio,
+                            "videoLength": video_length,
+                            "resolutionName": resolution_name,
+                        }
                     }
                 },
             },
         }
 
+    async def _generate_internal(
+        self, token: str, post_id: str, prompt: str,
+        aspect_ratio: str, video_length: int,
+        resolution_name: str, preset: str
+    ) -> AsyncGenerator[bytes, None]:
+        """内部生成逻辑"""
+        session = None
+        try:
+            headers = self._build_headers(token)
+            payload = self._build_payload(prompt, post_id, aspect_ratio, video_length, resolution_name, preset)
+
+            session = AsyncSession(impersonate=get_config("grok.browser"))
+            response = await session.post(
+                CHAT_API,
+                headers=headers,
+                data=orjson.dumps(payload),
+                timeout=self.timeout,
+                stream=True,
+                proxies=self._build_proxies(),
+            )
+
+            if response.status_code != 200:
+                logger.error(f"Video generation failed: {response.status_code}")
+                raise UpstreamException(
+                    message=f"Video generation failed: {response.status_code}",
+                    details={"status": response.status_code},
+                )
+
+            async def stream_response():
+                try:
+                    async for line in response.aiter_lines():
+                        yield line
+                finally:
+                    await session.close()
+
+            return stream_response()
+
+        except Exception as e:
+            if session:
+                try:
+                    await session.close()
+                except Exception:
+                    pass
+            logger.error(f"Video generation error: {e}")
+            if isinstance(e, AppException):
+                raise
+            raise UpstreamException(f"Video generation error: {str(e)}")
+
     async def generate(
-        self,
-        token: str,
-        prompt: str,
+        self, token: str, prompt: str,
         aspect_ratio: str = "3:2",
         video_length: int = 6,
         resolution_name: str = "480p",
-        preset: str = "normal",
+        preset: str = "normal"
     ) -> AsyncGenerator[bytes, None]:
-        """
-        生成视频
-
-        Args:
-            token: 认证 Token
-            prompt: 视频描述
-            aspect_ratio: 宽高比
-            video_length: 视频时长
-            resolution_name: 分辨率
-            preset: 预设
-
-        Returns:
-            AsyncGenerator，流式传输
-
-        Raises:
-            UpstreamException: 连接失败时
-        """
-        value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT)
-        try:
-            value = int(value)
-        except Exception:
-            value = DEFAULT_MAX_CONCURRENT
-        value = max(1, value)
-        global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE
-        if value != _MEDIA_SEM_VALUE:
-            _MEDIA_SEM_VALUE = value
-            _MEDIA_SEMAPHORE = asyncio.Semaphore(value)
-        async with _MEDIA_SEMAPHORE:
-            session = None
-            try:
-                # Step 1: 创建帖子
-                post_id = await self.create_post(token, prompt)
-
-                # Step 2: 建立连接
-                headers = self._build_headers(token)
-                payload = self._build_payload(
-                    prompt, post_id, aspect_ratio, video_length, resolution_name, preset
-                )
-
-                browser = get_config("grok.browser", DEFAULT_BROWSER)
-                session = AsyncSession(impersonate=browser)
-                response = await session.post(
-                    CHAT_API,
-                    headers=headers,
-                    data=orjson.dumps(payload),
-                    timeout=self.timeout,
-                    stream=True,
-                    proxies=self._build_proxies(),
-                )
-
-                if response.status_code != 200:
-                    logger.error(f"Video generation failed: {response.status_code}")
-                    try:
-                        await session.close()
-                    except Exception:
-                        pass
-                    raise UpstreamException(
-                        message=f"Video generation failed: {response.status_code}",
-                        details={"status": response.status_code},
-                    )
-
-                # Step 3: 流式传输
-                async def stream_response():
-                    try:
-                        async for line in response.aiter_lines():
-                            yield line
-                    finally:
-                        if session:
-                            await session.close()
-
-                return stream_response()
-
-            except Exception as e:
-                if session:
-                    try:
-                        await session.close()
-                    except Exception:
-                        pass
-                logger.error(f"Video generation error: {e}")
-                if isinstance(e, AppException):
-                    raise e
-                raise UpstreamException(f"Video generation error: {str(e)}")
+        """生成视频"""
+        async with _get_semaphore():
+            post_id = await self.create_post(token, prompt)
+            return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution_name, preset)
 
     async def generate_from_image(
-        self,
-        token: str,
-        prompt: str,
-        image_url: str,
+        self, token: str, prompt: str, image_url: str,
         aspect_ratio: str = "3:2",
         video_length: int = 6,
         resolution: str = "480p",
-        preset: str = "normal",
+        preset: str = "normal"
     ) -> AsyncGenerator[bytes, None]:
-        """
-        从图片生成视频
-
-        Args:
-            token: 认证 Token
-            prompt: 视频描述
-            image_url: 图片 URL
-            aspect_ratio: 宽高比
-            video_length: 视频时长
-            resolution: 分辨率
-            preset: 预设
-
-        Returns:
-            AsyncGenerator，流式传输
-        """
-        value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT)
-        try:
-            value = int(value)
-        except Exception:
-            value = DEFAULT_MAX_CONCURRENT
-        value = max(1, value)
-        global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE
-        if value != _MEDIA_SEM_VALUE:
-            _MEDIA_SEM_VALUE = value
-            _MEDIA_SEMAPHORE = asyncio.Semaphore(value)
-        async with _MEDIA_SEMAPHORE:
-            session = None
-            try:
-                # Step 1: 创建帖子
-                post_id = await self.create_image_post(token, image_url)
-
-                # Step 2: 建立连接
-                headers = self._build_headers(token)
-                payload = self._build_payload(
-                    prompt, post_id, aspect_ratio, video_length, resolution, preset
-                )
-
-                browser = get_config("grok.browser", DEFAULT_BROWSER)
-                session = AsyncSession(impersonate=browser)
-                response = await session.post(
-                    CHAT_API,
-                    headers=headers,
-                    data=orjson.dumps(payload),
-                    timeout=self.timeout,
-                    stream=True,
-                    proxies=self._build_proxies(),
-                )
-
-                if response.status_code != 200:
-                    logger.error(f"Video from image failed: {response.status_code}")
-                    try:
-                        await session.close()
-                    except Exception:
-                        pass
-                    raise UpstreamException(
-                        message=f"Video from image failed: {response.status_code}",
-                        details={"status": response.status_code},
-                    )
-
-                # Step 3: 流式传输
-                async def stream_response():
-                    try:
-                        async for line in response.aiter_lines():
-                            yield line
-                    finally:
-                        if session:
-                            await session.close()
-
-                return stream_response()
-
-            except Exception as e:
-                if session:
-                    try:
-                        await session.close()
-                    except Exception:
-                        pass
-                logger.error(f"Video from image error: {e}")
-                if isinstance(e, AppException):
-                    raise e
-                raise UpstreamException(f"Video from image error: {str(e)}")
-
-    @staticmethod
-    async def _wrap_stream(stream: AsyncGenerator, token_mgr, token: str, model: str):
-        """
-        包装流式响应，在完成时记录使用
-
-        Args:
-            stream: 原始 AsyncGenerator
-            token_mgr: TokenManager 实例
-            token: Token 字符串
-            model: 模型名称
-        """
-        success = False
-        try:
-            async for chunk in stream:
-                yield chunk
-            success = True
-        finally:
-            # 只在成功完成时记录使用，失败/异常时不扣费
-            if success:
-                try:
-                    model_info = ModelService.get(model)
-                    effort = (
-                        EffortType.HIGH
-                        if (model_info and model_info.cost.value == "high")
-                        else EffortType.LOW
-                    )
-                    await token_mgr.consume(token, effort)
-                    logger.debug(
-                        f"Video stream completed, recorded usage for token {token[:10]}... (effort={effort.value})"
-                    )
-                except Exception as e:
-                    logger.warning(f"Failed to record video stream usage: {e}")
+        """从图片生成视频"""
+        async with _get_semaphore():
+            post_id = await self.create_image_post(token, image_url)
+            return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution, preset)
 
     @staticmethod
     async def completions(
-        model: str,
-        messages: list,
+        model: str, messages: list,
         stream: bool = None,
         thinking: str = None,
         aspect_ratio: str = "3:2",
         video_length: int = 6,
         resolution: str = "480p",
-        preset: str = "normal",
+        preset: str = "normal"
     ):
-        """
-        视频生成入口
-
-        Args:
-            model: 模型名称
-            messages: 消息列表
-            stream: 是否流式
-            thinking: 思考模式
-            aspect_ratio: 宽高比
-            video_length: 视频时长
-            resolution: 分辨率
-            preset: 预设模式
-
-        Returns:
-            AsyncGenerator (流式) 或 dict (非流式)
-        """
+        """视频生成入口"""
         # 获取 token
-        try:
-            token_mgr = await get_token_manager()
-            await token_mgr.reload_if_stale()
-            token = None
-            for pool_name in ModelService.pool_candidates_for_model(model):
-                token = token_mgr.get_token(pool_name)
-                if token:
-                    break
-        except Exception as e:
-            logger.error(f"Failed to get token: {e}")
-            raise AppException(
-                message="Internal service error obtaining token",
-                error_type=ErrorType.SERVER.value,
-                code="internal_error",
-            )
+        token_mgr = await get_token_manager()
+        await token_mgr.reload_if_stale()
+        
+        token = None
+        for pool_name in ModelService.pool_candidates_for_model(model):
+            token = token_mgr.get_token(pool_name)
+            if token:
+                break
 
         if not token:
             raise AppException(
@@ -467,14 +262,8 @@ class VideoService:
                 status_code=429,
             )
 
-        # 解析参数
-        think = None
-        if thinking == "enabled":
-            think = True
-        elif thinking == "disabled":
-            think = False
-
-        is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM)
+        think = {"enabled": True, "disabled": False}.get(thinking)
+        is_stream = stream if stream is not None else get_config("grok.stream")
 
         # 提取内容
         from app.services.grok.services.chat import MessageExtractor
@@ -492,27 +281,18 @@ class VideoService:
             try:
                 for attach_type, attach_data in attachments:
                     if attach_type == "image":
-                        # 上传图片
                         _, file_uri = await upload_service.upload(attach_data, token)
                         image_url = f"https://assets.grok.com/{file_uri}"
                         logger.info(f"Image uploaded for video: {image_url}")
-                        break  # 视频模型只使用第一张图片
+                        break
             finally:
                 await upload_service.close()
 
         # 生成视频
         service = VideoService()
-
-        # 图片转视频
         if image_url:
             response = await service.generate_from_image(
-                token,
-                prompt,
-                image_url,
-                aspect_ratio,
-                video_length,
-                resolution,
-                preset,
+                token, prompt, image_url, aspect_ratio, video_length, resolution, preset
             )
         else:
             response = await service.generate(
@@ -522,26 +302,16 @@ class VideoService:
         # 处理响应
         if is_stream:
             processor = VideoStreamProcessor(model, token, think)
-            return VideoService._wrap_stream(
-                processor.process(response), token_mgr, token, model
-            )
-        else:
-            result = await VideoCollectProcessor(model, token).process(response)
-            # 非流式：处理完成后立即记录使用
-            try:
-                model_info = ModelService.get(model)
-                effort = (
-                    EffortType.HIGH
-                    if (model_info and model_info.cost.value == "high")
-                    else EffortType.LOW
-                )
-                await token_mgr.consume(token, effort)
-                logger.debug(
-                    f"Video completed, recorded usage for token {token[:10]}... (effort={effort.value})"
-                )
-            except Exception as e:
-                logger.warning(f"Failed to record video usage: {e}")
-            return result
-
+            return wrap_stream_with_usage(processor.process(response), token_mgr, token, model)
+        
+        result = await VideoCollectProcessor(model, token).process(response)
+        try:
+            model_info = ModelService.get(model)
+            effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW
+            await token_mgr.consume(token, effort)
+            logger.debug(f"Video completed, recorded usage (effort={effort.value})")
+        except Exception as e:
+            logger.warning(f"Failed to record video usage: {e}")
+        return result
 
 __all__ = ["VideoService"]

app/services/grok/services/nsfw.py

@@ -20,16 +20,9 @@ from app.services.grok.protocols.grpc_web import (
 )
 from app.services.grok.utils.headers import build_sso_cookie
 
-
 NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls"
 BIRTH_DATE_API = "https://grok.com/rest/auth/set-birth-date"
 
-DEFAULT_BROWSER = "chrome136"
-DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"
-DEFAULT_TIMEOUT = 30
-DEFAULT_BASE_PROXY_URL = ""
-
-
 @dataclass
 class NSFWResult:
     """NSFW 操作结果"""
@@ -40,20 +33,22 @@ class NSFWResult:
     grpc_message: Optional[str] = None
     error: Optional[str] = None
 
-
 class NSFWService:
     """NSFW 模式服务"""
 
     def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL)
-        self.timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT))
+        self.proxy = proxy or get_config("grok.base_proxy_url")
+        self.timeout = float(get_config("grok.timeout"))
+
+    def _build_proxies(self) -> Optional[dict]:
+        """构建代理配置"""
+        return {"http": self.proxy, "https": self.proxy} if self.proxy else None
 
     @staticmethod
     def _random_birth_date() -> str:
-        """生成随机出生日期"""
+        """生成随机出生日期（20-40岁）"""
         today = datetime.date.today()
-        age = random.randint(20, 40)
-        birth_year = today.year - age
+        birth_year = today.year - random.randint(20, 40)
         birth_month = random.randint(1, 12)
         birth_day = random.randint(1, 28)
         hour = random.randint(0, 23)
@@ -65,7 +60,7 @@ class NSFWService:
     def _build_headers(self, token: str) -> dict:
         """构造 gRPC-Web 请求头"""
         cookie = build_sso_cookie(token, include_rw=True)
-        user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT)
+        user_agent = get_config("grok.user_agent")
         return {
             "accept": "*/*",
             "content-type": "application/grpc-web+proto",
@@ -80,7 +75,7 @@ class NSFWService:
     def _build_birth_headers(self, token: str) -> dict:
         """构造设置出生日期请求头"""
         cookie = build_sso_cookie(token, include_rw=True)
-        user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT)
+        user_agent = get_config("grok.user_agent")
         return {
             "accept": "*/*",
             "content-type": "application/json",
@@ -102,18 +97,20 @@ class NSFWService:
         protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner
         return encode_grpc_web_payload(protobuf)
 
-    async def _set_birth_date(self, session: AsyncSession, token: str) -> tuple[bool, int, Optional[str]]:
+    async def _set_birth_date(
+        self, session: AsyncSession, token: str
+    ) -> tuple[bool, int, Optional[str]]:
         """设置出生日期"""
         headers = self._build_birth_headers(token)
         payload = {"birthDate": self._random_birth_date()}
-        proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None
+        
         try:
             response = await session.post(
                 BIRTH_DATE_API,
                 json=payload,
                 headers=headers,
                 timeout=self.timeout,
-                proxies=proxies,
+                proxies=self._build_proxies(),
             )
             if response.status_code in (200, 204):
                 return True, response.status_code, None
@@ -125,31 +122,27 @@ class NSFWService:
         """为单个 token 开启 NSFW 模式"""
         headers = self._build_headers(token)
         payload = self._build_payload()
-        logger.debug(
-            "NSFW payload: len={} hex={}",
-            len(payload),
-            payload.hex(),
-        )
-        proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None
+        logger.debug(f"NSFW payload: len={len(payload)} hex={payload.hex()}")
 
         try:
-            browser = get_config("grok.browser", DEFAULT_BROWSER)
+            browser = get_config("grok.browser")
             async with AsyncSession(impersonate=browser) as session:
-                ok, birth_status, birth_err = await self._set_birth_date(
-                    session, token
-                )
+                # 先设置出生日期
+                ok, birth_status, birth_err = await self._set_birth_date(session, token)
                 if not ok:
                     return NSFWResult(
                         success=False,
                         http_status=birth_status,
                         error=f"Set birth date failed: {birth_err}",
                     )
+                
+                # 开启 NSFW
                 response = await session.post(
                     NSFW_API,
                     data=payload,
                     headers=headers,
                     timeout=self.timeout,
-                    proxies=proxies,
+                    proxies=self._build_proxies(),
                 )
 
                 if response.status_code != 200:
@@ -160,18 +153,15 @@ class NSFWService:
                     )
 
                 # 解析 gRPC-Web 响应
-                content_type = response.headers.get("content-type")
                 _, trailers = parse_grpc_web_response(
-                    response.content, content_type=content_type
+                    response.content,
+                    content_type=response.headers.get("content-type")
                 )
 
                 grpc_status = get_grpc_status(trailers)
                 logger.debug(
-                    "NSFW response: http={} grpc={} msg={} trailers={}",
-                    response.status_code,
-                    grpc_status.code,
-                    grpc_status.message,
-                    trailers,
+                    f"NSFW response: http={response.status_code} grpc={grpc_status.code} "
+                    f"msg={grpc_status.message} trailers={trailers}"
                 )
 
                 # HTTP 200 且无 grpc-status（空响应）或 grpc-status=0 都算成功
@@ -188,5 +178,4 @@ class NSFWService:
             logger.error(f"NSFW enable failed: {e}")
             return NSFWResult(success=False, http_status=0, error=str(e)[:100])
 
-
 __all__ = ["NSFWService", "NSFWResult"]

app/services/grok/services/usage.py

@@ -13,32 +13,22 @@ from app.core.exceptions import UpstreamException
 from app.services.grok.utils.headers import apply_statsig, build_sso_cookie
 from app.services.grok.utils.retry import retry_on_status
 
-
 LIMITS_API = "https://grok.com/rest/rate-limits"
 
-DEFAULT_BROWSER = "chrome136"
-DEFAULT_USER_AGENT = (
-    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
-    "AppleWebKit/537.36 (KHTML, like Gecko) "
-    "Chrome/136.0.0.0 Safari/537.36"
-)
-DEFAULT_TIMEOUT = 10
-DEFAULT_MAX_CONCURRENT = 25
-DEFAULT_BASE_PROXY_URL = ""
-_USAGE_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT)
-_USAGE_SEM_VALUE = DEFAULT_MAX_CONCURRENT
+_USAGE_SEMAPHORE = asyncio.Semaphore(25)
+_USAGE_SEM_VALUE = 25
 
 
 class UsageService:
     """用量查询服务"""
 
     def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL)
-        self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT)
+        self.proxy = proxy or get_config("grok.base_proxy_url")
+        self.timeout = get_config("grok.timeout")
 
     def _build_headers(self, token: str) -> dict:
         """构建请求头"""
-        user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT)
+        user_agent = get_config("grok.user_agent")
         headers = {
             "Accept": "*/*",
             "Accept-Encoding": "gzip, deflate, br, zstd",
@@ -85,11 +75,11 @@ class UsageService:
         Raises:
             UpstreamException: 当获取失败且重试耗尽时
         """
-        value = get_config("performance.usage_max_concurrent", DEFAULT_MAX_CONCURRENT)
+        value = get_config("performance.usage_max_concurrent")
         try:
             value = int(value)
         except Exception:
-            value = DEFAULT_MAX_CONCURRENT
+            value = 25
         value = max(1, value)
         global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE
         if value != _USAGE_SEM_VALUE:
@@ -107,7 +97,7 @@ class UsageService:
                 try:
                     headers = self._build_headers(token)
                     payload = {"requestKind": "DEFAULT", "modelName": model_name}
-                    browser = get_config("grok.browser", DEFAULT_BROWSER)
+                    browser = get_config("grok.browser")
 
                     async with AsyncSession() as session:
                         response = await session.post(
@@ -152,5 +142,4 @@ class UsageService:
                 # 最后一次失败已经被记录
                 raise
 
-
 __all__ = ["UsageService"]

app/services/grok/services/voice.py

@@ -12,18 +12,13 @@ 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"
-DEFAULT_TIMEOUT = 30
-DEFAULT_BROWSER = "chrome136"
-DEFAULT_BASE_PROXY_URL = ""
-
 
 class VoiceService:
     """Voice Mode Service (LiveKit)"""
 
     def __init__(self, proxy: str = None):
-        self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL)
+        self.proxy = proxy or get_config("grok.base_proxy_url")
 
     async def get_token(
         self,
@@ -46,8 +41,8 @@ class VoiceService:
         proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None
         
         try:
-            browser = get_config("grok.browser", DEFAULT_BROWSER)
-            timeout = get_config("grok.timeout", DEFAULT_TIMEOUT)
+            browser = get_config("grok.browser")
+            timeout = get_config("grok.timeout")
             async with AsyncSession(impersonate=browser) as session:
                 response = await session.post(
                     LIVEKIT_TOKEN_API,

app/services/grok/utils/headers.py

@@ -11,16 +11,13 @@ from app.core.config import get_config
 from app.services.grok.utils.statsig import StatsigService
 
 
-DEFAULT_CF_CLEARANCE = ""
-
-
 def _normalize_token(token: str) -> str:
     return token[4:] if token.startswith("sso=") else token
 
 
 def build_sso_cookie(token: str, include_rw: bool = False) -> str:
     token = _normalize_token(token)
-    cf = get_config("grok.cf_clearance", DEFAULT_CF_CLEARANCE)
+    cf = get_config("grok.cf_clearance")
     cookie = f"sso={token}"
     if include_rw:
         cookie = f"{cookie}; sso-rw={token}"

app/services/grok/utils/retry.py

@@ -18,38 +18,22 @@ from app.core.config import get_config
 from app.core.exceptions import UpstreamException
 
 
-DEFAULT_MAX_RETRY = 3
-DEFAULT_RETRY_STATUS_CODES = [401, 429, 403]
-DEFAULT_RETRY_BACKOFF_BASE = 0.5
-DEFAULT_RETRY_BACKOFF_FACTOR = 2.0
-DEFAULT_RETRY_BACKOFF_MAX = 30.0
-DEFAULT_RETRY_BUDGET = 90.0
-
-
 class RetryContext:
     """重试上下文"""
 
     def __init__(self):
         self.attempt = 0
-        self.max_retry = int(get_config("grok.max_retry", DEFAULT_MAX_RETRY))
-        self.retry_codes = get_config(
-            "grok.retry_status_codes", DEFAULT_RETRY_STATUS_CODES
-        )
+        self.max_retry = int(get_config("grok.max_retry"))
+        self.retry_codes = get_config("grok.retry_status_codes")
         self.last_error = None
         self.last_status = None
         self.total_delay = 0.0
-        self.retry_budget = float(get_config("grok.retry_budget", DEFAULT_RETRY_BUDGET))
+        self.retry_budget = float(get_config("grok.retry_budget"))
 
         # 退避参数
-        self.backoff_base = float(
-            get_config("grok.retry_backoff_base", DEFAULT_RETRY_BACKOFF_BASE)
-        )
-        self.backoff_factor = float(
-            get_config("grok.retry_backoff_factor", DEFAULT_RETRY_BACKOFF_FACTOR)
-        )
-        self.backoff_max = float(
-            get_config("grok.retry_backoff_max", DEFAULT_RETRY_BACKOFF_MAX)
-        )
+        self.backoff_base = float(get_config("grok.retry_backoff_base"))
+        self.backoff_factor = float(get_config("grok.retry_backoff_factor"))
+        self.backoff_max = float(get_config("grok.retry_backoff_max"))
 
         # decorrelated jitter 状态
         self._last_delay = self.backoff_base

app/services/grok/utils/statsig.py

@@ -9,9 +9,6 @@ import string
 from app.core.config import get_config
 
 
-DEFAULT_DYNAMIC_STATSIG = True
-
-
 class StatsigService:
     """Statsig ID 生成服务"""
 
@@ -33,8 +30,7 @@ class StatsigService:
         Returns:
             Base64 编码的 ID
         """
-        # 读取配置
-        dynamic = get_config("grok.dynamic_statsig", DEFAULT_DYNAMIC_STATSIG)
+        dynamic = get_config("grok.dynamic_statsig")
 
         if not dynamic:
             return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk="

app/services/grok/utils/stream.py

@@ -0,0 +1,43 @@
+"""
+流式响应通用工具
+"""
+
+from typing import AsyncGenerator
+
+from app.core.logger import logger
+from app.services.grok.models.model import ModelService
+from app.services.token import EffortType
+
+
+async def wrap_stream_with_usage(
+    stream: AsyncGenerator, 
+    token_mgr, 
+    token: str, 
+    model: str
+) -> AsyncGenerator:
+    """
+    包装流式响应，在完成时记录使用
+    
+    Args:
+        stream: 原始 AsyncGenerator
+        token_mgr: TokenManager 实例
+        token: Token 字符串
+        model: 模型名称
+    """
+    success = False
+    try:
+        async for chunk in stream:
+            yield chunk
+        success = True
+    finally:
+        if success:
+            try:
+                model_info = ModelService.get(model)
+                effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW
+                await token_mgr.consume(token, effort)
+                logger.debug(f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})")
+            except Exception as e:
+                logger.warning(f"Failed to record stream usage: {e}")
+
+
+__all__ = ["wrap_stream_with_usage"]

data/config.toml

@@ -1,47 +1,36 @@
 [grok]
-temporary = true
-stream = true
-thinking = true
-dynamic_statsig = true
-filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"]
+browser = "chrome136"
 timeout = 120
+user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36"
 base_proxy_url = ""
 asset_proxy_url = ""
 cf_clearance = ""
-browser = "chrome136"
-user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36"
+temporary = true
 disable_memory = true
+stream = true
+thinking = true
+filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"]
 max_retry = 3
 retry_status_codes = [401,429,403]
-image_ws = true
-image_ws_blocked_seconds = 15
-image_ws_final_min_bytes = 100000
-image_ws_nsfw = true
 retry_backoff_base = 0.5
 retry_backoff_factor = 2.0
 retry_backoff_max = 30.0
 retry_budget = 90.0
 stream_idle_timeout = 120.0
 video_idle_timeout = 90.0
+image_ws_blocked_seconds = 15
+image_ws_final_min_bytes = 100000
+image_ws_medium_min_bytes = 30000
+dynamic_statsig = true
+image_ws = true
+image_ws_nsfw = true
 
 [app]
 app_url = "http://127.0.0.1:8000"
-app_key = "grok2api"
-api_key = ""
 image_format = "url"
 video_format = "html"
-
-[token]
-auto_refresh = true
-refresh_interval_hours = 8
-super_refresh_interval_hours = 2
-fail_threshold = 5
-save_delay_ms = 200
-reload_interval_sec = 30
-
-[cache]
-enable_auto_clean = true
-limit_mb = 1024
+app_key = "grok2api"
+api_key = ""
 
 [performance]
 assets_max_concurrent = 25
@@ -55,3 +44,15 @@ usage_max_tokens = 1000
 nsfw_max_concurrent = 10
 nsfw_batch_size = 50
 nsfw_max_tokens = 1000
+
+[cache]
+enable_auto_clean = true
+limit_mb = 1024
+
+[token]
+auto_refresh = true
+refresh_interval_hours = 8
+super_refresh_interval_hours = 2
+fail_threshold = 5
+save_delay_ms = 200
+reload_interval_sec = 30

main.py

@@ -41,17 +41,21 @@ setup_logging(
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """应用生命周期管理"""
-    # 1. 加载配置
-    from app.core.config import config
-
+    # 1. 注册服务默认配置
+    from app.core.config import config, register_defaults
+    from app.services.grok.defaults import get_grok_defaults
+    
+    register_defaults(get_grok_defaults())
+    
+    # 2. 加载配置
     await config.load()
 
-    # 2. 启动服务显示
+    # 3. 启动服务显示
     logger.info("Starting Grok2API...")
     logger.info(f"Platform: {platform.system()} {platform.release()}")
     logger.info(f"Python: {sys.version.split()[0]}")
 
-    # 3. 启动 Token 刷新调度器
+    # 4. 启动 Token 刷新调度器
     refresh_enabled = get_config("token.auto_refresh", True)
     if refresh_enabled:
         basic_interval = get_config("token.refresh_interval_hours", 8)

tests/test_model.py

@@ -18,7 +18,7 @@ load_dotenv(ROOT / ".env")
 from app.core.config import config
 from app.core.logger import setup_logging
 from app.services.grok.services.chat import GrokChatService
-from app.services.grok.processors.processor import CollectProcessor
+from app.services.grok.processors import CollectProcessor
 from app.services.grok.services.usage import UsageService