From 314102dcbadc910fbf449a4b94baa69ce989e89d Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Fri, 13 Feb 2026 13:00:12 +0800 Subject: [PATCH] feat: add NSFW configuration and refactor related services for improved management --- app/api/v1/admin.py | 29 +++-- app/core/logger.py | 4 + app/services/grok/batch_services/__init__.py | 4 +- app/services/grok/batch_services/assets.py | 7 +- app/services/grok/batch_services/nsfw.py | 129 +++++++++---------- app/services/grok/batch_services/usage.py | 6 +- app/services/grok/defaults.py | 7 +- app/services/grok/services/chat.py | 2 +- app/services/grok/services/video.py | 3 +- app/services/grok/services/voice.py | 2 +- app/services/grok/utils/__init__.py | 0 app/services/grok/utils/download.py | 2 +- app/services/grok/utils/upload.py | 2 +- app/services/reverse/accept_tos.py | 118 +++++++++++++++++ app/services/reverse/nsfw_mgmt.py | 12 +- app/services/reverse/rate_limits.py | 8 -- app/services/reverse/set_birth.py | 14 +- app/services/reverse/utils/grpc.py | 48 +++++++ app/services/reverse/utils/headers.py | 9 +- app/services/token/manager.py | 9 ++ app/services/token/service.py | 27 ++-- app/static/config/config.js | 12 +- config.defaults.toml | 14 +- data/config.toml | 27 +++- 24 files changed, 341 insertions(+), 154 deletions(-) delete mode 100644 app/services/grok/utils/__init__.py create mode 100644 app/services/reverse/accept_tos.py diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index ecd5caa..cc9c6fe 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -8,7 +8,7 @@ from fastapi import ( WebSocketDisconnect, ) from fastapi.responses import HTMLResponse, StreamingResponse, RedirectResponse -from typing import Optional +from typing import Optional, List, Tuple from pydantic import BaseModel from app.core.auth import verify_api_key, verify_app_key, get_admin_api_key from app.core.config import config, get_config @@ -16,7 +16,8 @@ from app.core.batch_tasks import create_task, get_task, expire_task from app.core.storage import get_storage, LocalStorage, RedisStorage, SQLStorage from app.core.exceptions import AppException from app.services.token.manager import get_token_manager -from app.services.grok.batch_services import BatchUsageService, BatchNSFWService +from app.services.grok.batch_services import BatchUsageService +from app.services.grok.batch_services.nsfw import NSFWService from app.services.grok.batch_services.assets import ListService, DeleteService import os import time @@ -87,7 +88,7 @@ async def _delete_imagine_session(task_id: str) -> None: _IMAGINE_SESSIONS.pop(task_id, None) -async def _delete_imagine_sessions(task_ids: list[str]) -> int: +async def _delete_imagine_sessions(task_ids: List[str]) -> int: if not task_ids: return 0 removed = 0 @@ -99,7 +100,7 @@ async def _delete_imagine_sessions(task_ids: list[str]) -> int: return removed -def _collect_tokens(data: dict) -> list[str]: +def _collect_tokens(data: dict) -> List[str]: """从请求数据中收集 token 列表""" tokens = [] if isinstance(data.get("token"), str) and data["token"].strip(): @@ -110,8 +111,8 @@ def _collect_tokens(data: dict) -> list[str]: def _truncate_tokens( - tokens: list[str], max_tokens: int, operation: str = "operation" -) -> tuple[list[str], bool, int]: + tokens: List[str], max_tokens: int, operation: str = "operation" +) -> Tuple[List[str], bool, int]: """去重并截断 token 列表，返回 (unique_tokens, truncated, original_count)""" unique_tokens = list(dict.fromkeys(tokens)) original_count = len(unique_tokens) @@ -538,7 +539,7 @@ async def admin_imagine_start(data: ImagineStartRequest): class ImagineStopRequest(BaseModel): - task_ids: list[str] + task_ids: List[str] @router.post("/api/v1/admin/imagine/stop", dependencies=[Depends(verify_api_key)]) @@ -970,10 +971,10 @@ async def enable_nsfw_api(data: dict): ) # 批量执行配置 - max_concurrent = get_config("performance.nsfw_max_concurrent") - batch_size = get_config("performance.nsfw_batch_size") + max_concurrent = get_config("nsfw.concurrent") + batch_size = get_config("nsfw.batch_size") - raw_results = await BatchNSFWService.enable( + raw_results = await NSFWService.batch( unique_tokens, mgr, max_concurrent=max_concurrent, @@ -1043,8 +1044,8 @@ async def enable_nsfw_api_async(data: dict): tokens, max_tokens, "NSFW enable" ) - max_concurrent = get_config("performance.nsfw_max_concurrent") - batch_size = get_config("performance.nsfw_batch_size") + max_concurrent = get_config("nsfw.concurrent") + batch_size = get_config("nsfw.batch_size") task = create_task(len(unique_tokens)) @@ -1055,7 +1056,7 @@ async def enable_nsfw_api_async(data: dict): ok = bool(res.get("ok") and res.get("data", {}).get("success")) task.record(ok) - raw_results = await BatchNSFWService.enable( + raw_results = await NSFWService.batch( unique_tokens, mgr, max_concurrent=max_concurrent, @@ -1328,7 +1329,7 @@ async def load_online_cache_api_async(data: dict): tokens = data.get("tokens") scope = data.get("scope") - selected_tokens: list[str] = [] + selected_tokens: List[str] = [] if isinstance(tokens, list): selected_tokens = [str(t).strip() for t in tokens if str(t).strip()] diff --git a/app/core/logger.py b/app/core/logger.py index a49b219..0b0290f 100644 --- a/app/core/logger.py +++ b/app/core/logger.py @@ -9,6 +9,10 @@ import traceback from pathlib import Path from loguru import logger +# Provide logging.Logger compatibility for legacy calls +if not hasattr(logger, "isEnabledFor"): + logger.isEnabledFor = lambda _level: True + # 日志目录 DEFAULT_LOG_DIR = Path(__file__).parent.parent.parent / "logs" LOG_DIR = Path(os.getenv("LOG_DIR", str(DEFAULT_LOG_DIR))) diff --git a/app/services/grok/batch_services/__init__.py b/app/services/grok/batch_services/__init__.py index 7a4107e..5866658 100644 --- a/app/services/grok/batch_services/__init__.py +++ b/app/services/grok/batch_services/__init__.py @@ -1,5 +1,5 @@ """Batch services.""" from .usage import BatchUsageService -from .nsfw import BatchNSFWService -__all__ = ["BatchUsageService", "BatchNSFWService"] + +__all__ = ["BatchUsageService"] diff --git a/app/services/grok/batch_services/assets.py b/app/services/grok/batch_services/assets.py index 14b2ccc..5ba9ab3 100644 --- a/app/services/grok/batch_services/assets.py +++ b/app/services/grok/batch_services/assets.py @@ -9,7 +9,8 @@ from curl_cffi.requests import AsyncSession from app.core.config import get_config from app.core.logger import logger -from app.services.reverse import AssetsListReverse, AssetsDeleteReverse +from app.services.reverse.assets_list import AssetsListReverse +from app.services.reverse.assets_delete import AssetsDeleteReverse from app.services.grok.utils.batch import run_in_batches @@ -102,7 +103,7 @@ class ListService(BaseAssetsService): @staticmethod async def fetch_assets_details( - tokens: list[str], + tokens: List[str], account_map: dict, *, max_concurrent: int, @@ -190,7 +191,7 @@ class DeleteService(BaseAssetsService): @staticmethod async def clear_assets( - tokens: list[str], + tokens: List[str], mgr, *, max_concurrent: int, diff --git a/app/services/grok/batch_services/nsfw.py b/app/services/grok/batch_services/nsfw.py index b701bc3..f53128b 100644 --- a/app/services/grok/batch_services/nsfw.py +++ b/app/services/grok/batch_services/nsfw.py @@ -2,7 +2,6 @@ Batch NSFW service. """ -from dataclasses import dataclass from typing import Callable, Awaitable, Dict, Any, Optional from curl_cffi.requests import AsyncSession @@ -10,66 +9,16 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config from app.core.exceptions import UpstreamException -from app.services.reverse import NsfwMgmtReverse, SetBirthReverse -from app.services.reverse.utils.grpc import GrpcStatus +from app.services.reverse.accept_tos import AcceptTosReverse +from app.services.reverse.nsfw_mgmt import NsfwMgmtReverse +from app.services.reverse.set_birth import SetBirthReverse from app.services.grok.utils.batch import run_in_batches -@dataclass -class NSFWResult: - """NSFW 操作结果""" - - success: bool - http_status: int - grpc_status: Optional[int] = None - grpc_message: Optional[str] = None - error: Optional[str] = None - - class NSFWService: """NSFW 模式服务""" - - async def enable(self, token: str) -> NSFWResult: - """为单个 token 开启 NSFW 模式""" - try: - browser = get_config("security.browser") - async with AsyncSession(impersonate=browser) as session: - # 先设置出生日期 - try: - await SetBirthReverse.request(session, token) - except UpstreamException as e: - status = None - if e.details and "status" in e.details: - status = e.details["status"] - else: - status = getattr(e, "status_code", None) - return NSFWResult( - success=False, - http_status=status or 0, - error=f"Set birth date failed: {str(e)}", - ) - - # 开启 NSFW - grpc_status: GrpcStatus = await NsfwMgmtReverse.request(session, token) - success = grpc_status.code in (-1, 0) - - return NSFWResult( - success=success, - http_status=200, - grpc_status=grpc_status.code, - grpc_message=grpc_status.message or None, - ) - - except Exception as e: - logger.error(f"NSFW enable failed: {e}") - return NSFWResult(success=False, http_status=0, error=str(e)[:100]) - - -class BatchNSFWService: - """Batch NSFW orchestration.""" - @staticmethod - async def enable( + async def batch( tokens: list[str], mgr, *, @@ -78,19 +27,63 @@ class BatchNSFWService: on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, should_cancel: Optional[Callable[[], bool]] = None, ) -> Dict[str, Dict[str, Any]]: - nsfw_service = NSFWService() - + """Batch enable NSFW.""" async def _enable(token: str): - result = await nsfw_service.enable(token) - if result.success: - await mgr.add_tag(token, "nsfw") - return { - "success": result.success, - "http_status": result.http_status, - "grpc_status": result.grpc_status, - "grpc_message": result.grpc_message, - "error": result.error, - } + try: + browser = get_config("security.browser") + async with AsyncSession(impersonate=browser) as session: + async def _record_fail(err: UpstreamException, reason: str): + status = None + if err.details and "status" in err.details: + status = err.details["status"] + else: + status = getattr(err, "status_code", None) + if status in (401, 403): + await mgr.record_fail(token, status, reason) + return status or 0 + + try: + await AcceptTosReverse.request(session, token) + except UpstreamException as e: + status = await _record_fail(e, "tos_auth_failed") + return { + "success": False, + "http_status": status, + "error": f"Accept ToS failed: {str(e)}", + } + + try: + await SetBirthReverse.request(session, token) + except UpstreamException as e: + status = await _record_fail(e, "set_birth_auth_failed") + return { + "success": False, + "http_status": status, + "error": f"Set birth date failed: {str(e)}", + } + + try: + grpc_status = await NsfwMgmtReverse.request(session, token) + success = grpc_status.code in (-1, 0) + except UpstreamException as e: + status = await _record_fail(e, "nsfw_mgmt_auth_failed") + return { + "success": False, + "http_status": status, + "error": f"NSFW enable failed: {str(e)}", + } + if success: + await mgr.add_tag(token, "nsfw") + return { + "success": success, + "http_status": 200, + "grpc_status": grpc_status.code, + "grpc_message": grpc_status.message or None, + "error": None, + } + except Exception as e: + logger.error(f"NSFW enable failed: {e}") + return {"success": False, "http_status": 0, "error": str(e)[:100]} return await run_in_batches( tokens, @@ -102,4 +95,4 @@ class BatchNSFWService: ) -__all__ = ["BatchNSFWService", "NSFWService", "NSFWResult"] +__all__ = ["NSFWService"] diff --git a/app/services/grok/batch_services/usage.py b/app/services/grok/batch_services/usage.py index edac6f4..3bf66ba 100644 --- a/app/services/grok/batch_services/usage.py +++ b/app/services/grok/batch_services/usage.py @@ -3,13 +3,13 @@ Batch usage service. """ import asyncio -from typing import Callable, Awaitable, Dict, Any, Optional +from typing import Callable, Awaitable, Dict, Any, Optional, List from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config -from app.services.reverse import RateLimitsReverse +from app.services.reverse.rate_limits import RateLimitsReverse from app.services.grok.utils.batch import run_in_batches _USAGE_SEMAPHORE = asyncio.Semaphore(25) @@ -63,7 +63,7 @@ class BatchUsageService: @staticmethod async def refresh( - tokens: list[str], + tokens: List[str], mgr, *, max_concurrent: int, diff --git a/app/services/grok/defaults.py b/app/services/grok/defaults.py index b273dcc..4307fa2 100644 --- a/app/services/grok/defaults.py +++ b/app/services/grok/defaults.py @@ -74,13 +74,16 @@ GROK_DEFAULTS = { "delete_timeout": 60, "delete_batch_size": 10, }, + "nsfw": { + "concurrent": 10, + "batch_size": 50, + "timeout": 60, + }, "performance": { "media_max_concurrent": 50, "usage_max_concurrent": 25, "usage_batch_size": 50, "usage_max_tokens": 1000, - "nsfw_max_concurrent": 10, - "nsfw_batch_size": 50, "nsfw_max_tokens": 1000, }, } diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index aa43923..0dd27ce 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -18,7 +18,7 @@ from app.core.exceptions import ( from app.services.grok.models.model import ModelService from app.services.grok.utils.upload import UploadService from app.services.grok.processors import StreamProcessor, CollectProcessor -from app.services.reverse import AppChatReverse +from app.services.reverse.app_chat import AppChatReverse from app.services.grok.utils.stream import wrap_stream_with_usage from app.services.token import get_token_manager, EffortType diff --git a/app/services/grok/services/video.py b/app/services/grok/services/video.py index cd65c8a..56ff592 100644 --- a/app/services/grok/services/video.py +++ b/app/services/grok/services/video.py @@ -18,7 +18,8 @@ from app.services.grok.models.model import ModelService from app.services.token import get_token_manager, EffortType from app.services.grok.processors import VideoStreamProcessor, VideoCollectProcessor from app.services.grok.utils.stream import wrap_stream_with_usage -from app.services.reverse import AppChatReverse, MediaPostReverse +from app.services.reverse.app_chat import AppChatReverse +from app.services.reverse.media_post import MediaPostReverse _MEDIA_SEMAPHORE = None _MEDIA_SEM_VALUE = 0 diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 81515dc..a6dc5ed 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -7,7 +7,7 @@ from typing import Any, Dict from curl_cffi.requests import AsyncSession from app.core.config import get_config -from app.services.reverse import LivekitTokenReverse +from app.services.reverse.ws_livekit import LivekitTokenReverse class VoiceService: diff --git a/app/services/grok/utils/__init__.py b/app/services/grok/utils/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/app/services/grok/utils/download.py b/app/services/grok/utils/download.py index fbb7872..a13c717 100644 --- a/app/services/grok/utils/download.py +++ b/app/services/grok/utils/download.py @@ -19,7 +19,7 @@ from app.core.logger import logger from app.core.storage import DATA_DIR from app.core.config import get_config from app.core.exceptions import AppException -from app.services.reverse import AssetsDownloadReverse +from app.services.reverse.assets_download import AssetsDownloadReverse from app.services.grok.utils.locks import _get_download_semaphore, _file_lock diff --git a/app/services/grok/utils/upload.py b/app/services/grok/utils/upload.py index ee17cc4..9670792 100644 --- a/app/services/grok/utils/upload.py +++ b/app/services/grok/utils/upload.py @@ -19,7 +19,7 @@ from app.core.config import get_config from app.core.exceptions import AppException, UpstreamException, ValidationException from app.core.logger import logger from app.core.storage import DATA_DIR -from app.services.reverse import AssetsUploadReverse +from app.services.reverse.assets_upload import AssetsUploadReverse from app.services.grok.utils.locks import _get_upload_semaphore, _file_lock diff --git a/app/services/reverse/accept_tos.py b/app/services/reverse/accept_tos.py new file mode 100644 index 0000000..203e1f6 --- /dev/null +++ b/app/services/reverse/accept_tos.py @@ -0,0 +1,118 @@ +""" +Reverse interface: accept ToS (gRPC-Web). +""" + +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException +from app.services.reverse.utils.headers import build_headers +from app.services.reverse.utils.retry import retry_on_status +from app.services.reverse.utils.grpc import GrpcClient, GrpcStatus + +ACCEPT_TOS_API = "https://accounts.x.ai/auth_mgmt.AuthManagement/SetTosAcceptedVersion" + + +class AcceptTosReverse: + """/auth_mgmt.AuthManagement/SetTosAcceptedVersion reverse interface.""" + + @staticmethod + async def request(session: AsyncSession, token: str) -> GrpcStatus: + """Accept ToS via gRPC-Web. + + Args: + session: AsyncSession, the session to use for the request. + token: str, the SSO token. + + Returns: + GrpcStatus: Parsed gRPC status. + """ + try: + # Get proxies + base_proxy = get_config("network.base_proxy_url") + proxies = {"http": base_proxy, "https": base_proxy} if base_proxy else None + + # Build headers + headers = build_headers( + cookie_token=token, + origin="https://accounts.x.ai", + referer="https://accounts.x.ai/accept-tos", + ) + headers["Content-Type"] = "application/grpc-web+proto" + headers["Accept"] = "*/*" + headers["Sec-Fetch-Dest"] = "empty" + headers["x-grpc-web"] = "1" + headers["x-user-agent"] = "connect-es/2.1.1" + headers["Cache-Control"] = "no-cache" + headers["Pragma"] = "no-cache" + + # Build payload + payload = GrpcClient.encode_payload(b"\x10\x01") + + # Curl Config + timeout = get_config("nsfw.timeout") + browser = get_config("security.browser") + + async def _do_request(): + response = await session.post( + ACCEPT_TOS_API, + headers=headers, + data=payload, + timeout=timeout, + proxies=proxies, + impersonate=browser, + ) + + if response.status_code != 200: + logger.error( + f"AcceptTosReverse: Request failed, {response.status_code}", + extra={"error_type": "UpstreamException"}, + ) + raise UpstreamException( + message=f"AcceptTosReverse: Request failed, {response.status_code}", + details={"status": response.status_code}, + ) + + logger.debug(f"AcceptTosReverse: Request successful, {response.status_code}") + + return response + + response = await retry_on_status(_do_request) + + _, trailers = GrpcClient.parse_response( + response.content, + content_type=response.headers.get("content-type"), + headers=response.headers, + ) + grpc_status = GrpcClient.get_status(trailers) + + if grpc_status.code not in (-1, 0): + raise UpstreamException( + message=f"AcceptTosReverse: gRPC failed, {grpc_status.code}", + details={ + "status": grpc_status.http_equiv, + "grpc_status": grpc_status.code, + "grpc_message": grpc_status.message, + }, + ) + + return grpc_status + + except Exception as e: + # Handle upstream exception + if isinstance(e, UpstreamException): + raise + + # Handle other non-upstream exceptions + logger.error( + f"AcceptTosReverse: Request failed, {str(e)}", + extra={"error_type": type(e).__name__}, + ) + raise UpstreamException( + message=f"AcceptTosReverse: Request failed, {str(e)}", + details={"status": 502, "error": str(e)}, + ) + + +__all__ = ["AcceptTosReverse"] diff --git a/app/services/reverse/nsfw_mgmt.py b/app/services/reverse/nsfw_mgmt.py index 349e241..8056b23 100644 --- a/app/services/reverse/nsfw_mgmt.py +++ b/app/services/reverse/nsfw_mgmt.py @@ -7,7 +7,6 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config from app.core.exceptions import UpstreamException -from app.services.token.service import TokenService from app.services.reverse.utils.headers import build_headers from app.services.reverse.utils.retry import retry_on_status from app.services.reverse.utils.grpc import GrpcClient, GrpcStatus @@ -55,7 +54,7 @@ class NsfwMgmtReverse: payload = GrpcClient.encode_payload(protobuf) # Curl Config - timeout = get_config("network.timeout") + timeout = get_config("nsfw.timeout") browser = get_config("security.browser") async def _do_request(): @@ -78,6 +77,8 @@ class NsfwMgmtReverse: details={"status": response.status_code}, ) + logger.debug(f"NsfwMgmtReverse: Request successful, {response.status_code}") + return response response = await retry_on_status(_do_request) @@ -109,13 +110,6 @@ class NsfwMgmtReverse: status = e.details["status"] else: status = getattr(e, "status_code", None) - if status == 401: - try: - await TokenService.record_fail( - token, status, "nsfw_mgmt_auth_failed" - ) - except Exception: - pass raise # Handle other non-upstream exceptions diff --git a/app/services/reverse/rate_limits.py b/app/services/reverse/rate_limits.py index aa85260..5efd3c0 100644 --- a/app/services/reverse/rate_limits.py +++ b/app/services/reverse/rate_limits.py @@ -9,7 +9,6 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config from app.core.exceptions import UpstreamException -from app.services.token.service import TokenService from app.services.reverse.utils.headers import build_headers from app.services.reverse.utils.retry import retry_on_status @@ -85,13 +84,6 @@ class RateLimitsReverse: status = e.details["status"] else: status = getattr(e, "status_code", None) - if status == 401: - try: - await TokenService.record_fail( - token, status, "rate_limits_auth_failed" - ) - except Exception: - pass raise # Handle other non-upstream exceptions diff --git a/app/services/reverse/set_birth.py b/app/services/reverse/set_birth.py index f1b7221..556a331 100644 --- a/app/services/reverse/set_birth.py +++ b/app/services/reverse/set_birth.py @@ -10,7 +10,6 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config from app.core.exceptions import UpstreamException -from app.services.token.service import TokenService from app.services.reverse.utils.headers import build_headers from app.services.reverse.utils.retry import retry_on_status @@ -46,7 +45,7 @@ class SetBirthReverse: # Build payload today = datetime.date.today() - birth_year = today.year - random.randint(20, 40) + birth_year = today.year - random.randint(20, 48) birth_month = random.randint(1, 12) birth_day = random.randint(1, 28) hour = random.randint(0, 23) @@ -59,7 +58,7 @@ class SetBirthReverse: } # Curl Config - timeout = get_config("network.timeout") + timeout = get_config("nsfw.timeout") browser = get_config("security.browser") async def _do_request(): @@ -82,6 +81,8 @@ class SetBirthReverse: details={"status": response.status_code}, ) + logger.debug(f"SetBirthReverse: Request successful, {response.status_code}") + return response return await retry_on_status(_do_request) @@ -94,13 +95,6 @@ class SetBirthReverse: status = e.details["status"] else: status = getattr(e, "status_code", None) - if status == 401: - try: - await TokenService.record_fail( - token, status, "set_birth_auth_failed" - ) - except Exception: - pass raise # Handle other non-upstream exceptions diff --git a/app/services/reverse/utils/grpc.py b/app/services/reverse/utils/grpc.py index 446ba4b..39eb678 100644 --- a/app/services/reverse/utils/grpc.py +++ b/app/services/reverse/utils/grpc.py @@ -3,12 +3,15 @@ gRPC-Web helpers for reverse interfaces. """ import base64 +import json import re import struct from dataclasses import dataclass from typing import Dict, List, Mapping, Optional, Tuple from urllib.parse import unquote +from app.core.logger import logger + # Base64 正则 B64_RE = re.compile(rb"^[A-Za-z0-9+/=\r\n]+$") @@ -38,6 +41,22 @@ class GrpcStatus: class GrpcClient: """gRPC-Web helpers wrapper.""" + @staticmethod + def _safe_headers(headers: Optional[Mapping[str, str]]) -> Dict[str, str]: + if not headers: + return {} + safe: Dict[str, str] = {} + for k, v in headers.items(): + if k.lower() in ("set-cookie", "cookie", "authorization"): + safe[k] = "" + else: + safe[k] = str(v) + return safe + + @staticmethod + def _b64(data: bytes) -> str: + return base64.b64encode(data).decode() + @staticmethod def encode_payload(data: bytes) -> bytes: """Encode gRPC-Web data frame.""" @@ -118,6 +137,35 @@ class GrpcClient: if "grpc-message" in lower and "grpc-message" not in trailers: trailers["grpc-message"] = unquote(str(lower["grpc-message"]).strip()) + # Log full response details on gRPC error + raw_status = str(trailers.get("grpc-status", "")).strip() + try: + status_code = int(raw_status) + except Exception: + status_code = -1 + + if status_code not in (0, -1): + try: + payload = { + "grpc_status": status_code, + "grpc_message": trailers.get("grpc-message", ""), + "content_type": content_type or "", + "headers": cls._safe_headers(headers), + "trailers": trailers, + "messages_b64": [cls._b64(m) for m in messages], + "body_b64": cls._b64(body), + } + logger.error( + "gRPC response error: {}", + json.dumps(payload, ensure_ascii=False), + extra={"error_type": "GrpcError"}, + ) + except Exception as e: + logger.error( + f"gRPC response error: failed to log payload ({e})", + extra={"error_type": "GrpcError"}, + ) + return messages, trailers @staticmethod diff --git a/app/services/reverse/utils/headers.py b/app/services/reverse/utils/headers.py index 4388bf2..03a8f25 100644 --- a/app/services/reverse/utils/headers.py +++ b/app/services/reverse/utils/headers.py @@ -123,11 +123,10 @@ def build_headers(cookie_token: str, content_type: Optional[str] = None, origin: headers["x-xai-request-id"] = str(uuid.uuid4()) # Print headers without Cookie - if logger.isEnabledFor(10): - safe_headers = dict(headers) - if "Cookie" in safe_headers: - safe_headers["Cookie"] = "" - logger.debug(f"Built headers: {orjson.dumps(safe_headers).decode()}") + safe_headers = dict(headers) + if "Cookie" in safe_headers: + safe_headers["Cookie"] = "" + logger.debug(f"Built headers: {orjson.dumps(safe_headers).decode()}") return headers diff --git a/app/services/token/manager.py b/app/services/token/manager.py index 4f056d8..68bfd84 100644 --- a/app/services/token/manager.py +++ b/app/services/token/manager.py @@ -16,6 +16,7 @@ from app.services.token.models import ( ) from app.core.storage import get_storage, LocalStorage from app.core.config import get_config +from app.core.exceptions import UpstreamException from app.services.token.pool import TokenPool from app.services.grok.batch_services.usage import UsageService @@ -382,6 +383,14 @@ class TokenManager: return True except Exception as e: + if isinstance(e, UpstreamException): + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + if status in (401, 403): + await self.record_fail(token_str, status, "rate_limits_auth_failed") logger.warning( f"Token {raw_token[:10]}...: API sync failed, fallback to local ({e})" ) diff --git a/app/services/token/service.py b/app/services/token/service.py index 63e635c..b441fbe 100644 --- a/app/services/token/service.py +++ b/app/services/token/service.py @@ -2,7 +2,6 @@ from typing import List, Optional, Dict -from app.services.token.manager import get_token_manager from app.services.token.models import TokenInfo, EffortType @@ -13,6 +12,12 @@ class TokenService: 提供简化的 API，隐藏内部实现细节 """ + @staticmethod + async def _get_manager(): + from app.services.token.manager import get_token_manager + + return await get_token_manager() + @staticmethod async def get_token(pool_name: str = "ssoBasic") -> Optional[str]: """ @@ -24,7 +29,7 @@ class TokenService: Returns: Token 字符串（不含 sso= 前缀）或 None """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return manager.get_token(pool_name) @staticmethod @@ -39,7 +44,7 @@ class TokenService: Returns: 是否成功 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return await manager.consume(token, effort) @staticmethod @@ -54,7 +59,7 @@ class TokenService: Returns: 是否成功 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return await manager.sync_usage(token, effort) @staticmethod @@ -70,7 +75,7 @@ class TokenService: Returns: 是否成功 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return await manager.record_fail(token, status_code, reason) @staticmethod @@ -85,7 +90,7 @@ class TokenService: Returns: 是否成功 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return await manager.add(token, pool_name) @staticmethod @@ -99,7 +104,7 @@ class TokenService: Returns: 是否成功 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return await manager.remove(token) @staticmethod @@ -113,13 +118,13 @@ class TokenService: Returns: 是否成功 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return await manager.reset_token(token) @staticmethod async def reset_all(): """重置所有 Token""" - manager = await get_token_manager() + manager = await TokenService._get_manager() await manager.reset_all() @staticmethod @@ -130,7 +135,7 @@ class TokenService: Returns: 各池的统计信息 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return manager.get_stats() @staticmethod @@ -144,7 +149,7 @@ class TokenService: Returns: Token 列表 """ - manager = await get_token_manager() + manager = await TokenService._get_manager() return manager.get_pool_tokens(pool_name) diff --git a/app/static/config/config.js b/app/static/config/config.js index ee60d64..815a177 100644 --- a/app/static/config/config.js +++ b/app/static/config/config.js @@ -33,8 +33,8 @@ const NUMERIC_FIELDS = new Set([ 'image_ws_blocked_seconds', 'image_ws_final_min_bytes', 'image_ws_medium_min_bytes', - 'nsfw_max_concurrent', - 'nsfw_batch_size', + 'concurrent', + 'batch_size', 'nsfw_max_tokens' ]); @@ -117,11 +117,15 @@ const LOCALE_MAP = { "delete_timeout": { title: "删除超时", desc: "资产删除接口超时时间（秒）。推荐 60。" }, "delete_batch_size": { title: "删除批次大小", desc: "单次删除可处理的 Token 数量。推荐 10。" } }, + "nsfw": { + "label": "NSFW 配置", + "concurrent": { title: "并发上限", desc: "批量开启 NSFW 模式时的并发请求上限。推荐 10。" }, + "batch_size": { title: "批次大小", desc: "批量开启 NSFW 模式的单批处理数量。推荐 50。" }, + "timeout": { title: "请求超时", desc: "NSFW 开启相关请求的超时时间（秒）。推荐 60。" } + }, "performance": { "label": "并发性能", "media_max_concurrent": { title: "Media 并发上限", desc: "视频/媒体生成请求的并发上限。推荐 50。" }, - "nsfw_max_concurrent": { title: "NSFW 开启并发上限", desc: "批量开启 NSFW 模式时的并发请求上限。推荐 10。" }, - "nsfw_batch_size": { title: "NSFW 开启批量大小", desc: "批量开启 NSFW 模式的单批处理数量。推荐 50。" }, "nsfw_max_tokens": { title: "NSFW 开启最大数量", desc: "单次批量开启 NSFW 的 Token 数量上限，防止误操作。推荐 1000。" }, "usage_max_concurrent": { title: "Token 刷新并发上限", desc: "批量刷新 Token 用量时的并发请求上限。推荐 25。" }, "usage_batch_size": { title: "Token 刷新批次大小", desc: "批量刷新 Token 用量的单批处理数量。推荐 50。" }, diff --git a/config.defaults.toml b/config.defaults.toml index e15bed8..7b15856 100644 --- a/config.defaults.toml +++ b/config.defaults.toml @@ -132,6 +132,16 @@ delete_timeout = 60 # 资产删除批次大小（Token 维度） delete_batch_size = 10 +# ==================== NSFW ==================== +[nsfw] + +# NSFW 批量开启并发上限 +concurrent = 10 +# NSFW 批量开启批次大小 +batch_size = 50 +# NSFW 请求超时时间（秒） +timeout = 60 + # ==================== 并发性能 ==================== [performance] @@ -146,9 +156,5 @@ usage_batch_size = 50 # Token 用量单次最大数量 usage_max_tokens = 1000 -# NSFW 批量开启并发上限 -nsfw_max_concurrent = 10 -# NSFW 批量开启批次大小 -nsfw_batch_size = 50 # NSFW 单次最大数量 nsfw_max_tokens = 1000 diff --git a/data/config.toml b/data/config.toml index 93016ea..f5face4 100644 --- a/data/config.toml +++ b/data/config.toml @@ -54,15 +54,30 @@ reload_interval_sec = 30 enable_auto_clean = true limit_mb = 1024 +[asset] +upload_concurrent = 30 +upload_timeout = 60 +download_concurrent = 30 +download_timeout = 60 +list_concurrent = 10 +list_timeout = 60 +list_batch_size = 10 +delete_concurrent = 10 +delete_timeout = 60 +delete_batch_size = 10 + +[nsfw] +concurrent = 10 +batch_size = 50 +timeout = 60 + [performance] -assets_max_concurrent = 25 -assets_delete_batch_size = 10 -assets_batch_size = 10 -assets_max_tokens = 1000 media_max_concurrent = 50 usage_max_concurrent = 25 usage_batch_size = 50 usage_max_tokens = 1000 -nsfw_max_concurrent = 10 -nsfw_batch_size = 50 nsfw_max_tokens = 1000 +assets_max_concurrent = 25 +assets_delete_batch_size = 10 +assets_batch_size = 10 +assets_max_tokens = 1000