feat: 新增自定义解析接口、自动设置用户名、视频额度刷新

src/services/generation_handler.py

@@ -92,13 +92,13 @@ class GenerationHandler:
         is_video = model_config["type"] == "video"
         is_image = model_config["type"] == "image"
 
-        # Select token (with lock for image generation)
-        token_obj = await self.load_balancer.select_token(for_image_generation=is_image)
+        # Select token (with lock for image generation, Sora2 quota check for video generation)
+        token_obj = await self.load_balancer.select_token(for_image_generation=is_image, for_video_generation=is_video)
         if not token_obj:
             if is_image:
                 raise Exception("No available tokens for image generation. All tokens are either disabled, cooling down, locked, or expired.")
             else:
-                raise Exception("No available tokens. All tokens are either disabled, cooling down, or expired.")
+                raise Exception("No available tokens for video generation. All tokens are either disabled, cooling down, Sora2 quota exhausted, don't support Sora2, or expired.")
 
         # Acquire lock for image generation
         if is_image:
@@ -180,11 +180,7 @@ class GenerationHandler:
                 yield chunk
             
             # Record success
-            await self.token_manager.record_success(token_obj.id)
-
-            # Check cooldown for video
-            if is_video:
-                await self.token_manager.check_and_apply_cooldown(token_obj.id)
+            await self.token_manager.record_success(token_obj.id, is_video=is_video)
 
             # Release lock for image generation
             if is_image:
@@ -231,6 +227,8 @@ class GenerationHandler:
         max_attempts = int(timeout / poll_interval)  # Calculate max attempts based on timeout
         last_progress = 0
         start_time = time.time()
+        last_heartbeat_time = start_time  # Track last heartbeat for image generation
+        heartbeat_interval = 10  # Send heartbeat every 10 seconds for image generation
 
         debug_logger.log_info(f"Starting task polling: task_id={task_id}, is_video={is_video}, timeout={timeout}s, max_attempts={max_attempts}")
 
@@ -315,6 +313,10 @@ class GenerationHandler:
                                             reasoning_content="**Video Generation Completed**\n\nWatermark-free mode enabled. Publishing video to get watermark-free version...\n"
                                         )
 
+                                    # Get watermark-free config to determine parse method
+                                    watermark_config = await self.db.get_watermark_free_config()
+                                    parse_method = watermark_config.parse_method or "third_party"
+
                                     # Post video to get watermark-free version
                                     try:
                                         debug_logger.log_info(f"Calling post_video_for_watermark_free with generation_id={generation_id}, prompt={prompt[:50]}...")
@@ -328,8 +330,28 @@ class GenerationHandler:
                                         if not post_id:
                                             raise Exception("Failed to get post ID from publish API")
 
-                                        # Construct watermark-free video URL
-                                        watermark_free_url = f"https://oscdn2.dyysy.com/MP4/{post_id}.mp4"
+                                        # Get watermark-free video URL based on parse method
+                                        if parse_method == "custom":
+                                            # Use custom parse server
+                                            if not watermark_config.custom_parse_url or not watermark_config.custom_parse_token:
+                                                raise Exception("Custom parse server URL or token not configured")
+
+                                            if stream:
+                                                yield self._format_stream_chunk(
+                                                    reasoning_content=f"Video published successfully. Post ID: {post_id}\nUsing custom parse server to get watermark-free URL...\n"
+                                                )
+
+                                            debug_logger.log_info(f"Using custom parse server: {watermark_config.custom_parse_url}")
+                                            watermark_free_url = await self.sora_client.get_watermark_free_url_custom(
+                                                parse_url=watermark_config.custom_parse_url,
+                                                parse_token=watermark_config.custom_parse_token,
+                                                post_id=post_id
+                                            )
+                                        else:
+                                            # Use third-party parse (default)
+                                            watermark_free_url = f"https://oscdn2.dyysy.com/MP4/{post_id}.mp4"
+                                            debug_logger.log_info(f"Using third-party parse server")
+
                                         debug_logger.log_info(f"Watermark-free URL: {watermark_free_url}")
 
                                         if stream:
@@ -439,8 +461,10 @@ class GenerationHandler:
                     task_responses = result.get("task_responses", [])
 
                     # Find matching task
+                    task_found = False
                     for task_resp in task_responses:
                         if task_resp.get("id") == task_id:
+                            task_found = True
                             status = task_resp.get("status")
                             progress = task_resp.get("progress_pct", 0) * 100
 
@@ -513,6 +537,26 @@ class GenerationHandler:
                                             reasoning_content=f"**Processing**\n\nGeneration in progress: {progress:.0f}% completed...\n"
                                         )
 
+                    # For image generation, send heartbeat every 10 seconds if no progress update
+                    if not is_video and stream:
+                        current_time = time.time()
+                        if current_time - last_heartbeat_time >= heartbeat_interval:
+                            last_heartbeat_time = current_time
+                            elapsed = int(current_time - start_time)
+                            yield self._format_stream_chunk(
+                                reasoning_content=f"**Generating**\n\nImage generation in progress... ({elapsed}s elapsed)\n"
+                            )
+
+                    # If task not found in response, send heartbeat for image generation
+                    if not task_found and not is_video and stream:
+                        current_time = time.time()
+                        if current_time - last_heartbeat_time >= heartbeat_interval:
+                            last_heartbeat_time = current_time
+                            elapsed = int(current_time - start_time)
+                            yield self._format_stream_chunk(
+                                reasoning_content=f"**Generating**\n\nImage generation in progress... ({elapsed}s elapsed)\n"
+                            )
+
                 # Progress update for stream mode (fallback if no status from API)
                 if stream and attempt % 10 == 0:  # Update every 10 attempts (roughly 20% intervals)
                     estimated_progress = min(90, (attempt / max_attempts) * 100)

src/services/load_balancer.py

@@ -14,12 +14,13 @@ class LoadBalancer:
         # Use image timeout from config as lock timeout
         self.token_lock = TokenLock(lock_timeout=config.image_timeout)
 
-    async def select_token(self, for_image_generation: bool = False) -> Optional[Token]:
+    async def select_token(self, for_image_generation: bool = False, for_video_generation: bool = False) -> Optional[Token]:
         """
         Select a token using random load balancing
 
         Args:
             for_image_generation: If True, only select tokens that are not locked for image generation
+            for_video_generation: If True, filter out tokens with Sora2 quota exhausted (sora2_cooldown_until not expired) and tokens that don't support Sora2
 
         Returns:
             Selected token or None if no available tokens
@@ -29,6 +30,33 @@ class LoadBalancer:
         if not active_tokens:
             return None
 
+        # If for video generation, filter out tokens with Sora2 quota exhausted and tokens without Sora2 support
+        if for_video_generation:
+            from datetime import datetime
+            available_tokens = []
+            for token in active_tokens:
+                # Skip tokens that don't support Sora2
+                if not token.sora2_supported:
+                    continue
+
+                # Check if Sora2 cooldown has expired and refresh if needed
+                if token.sora2_cooldown_until and token.sora2_cooldown_until <= datetime.now():
+                    await self.token_manager.refresh_sora2_remaining_if_cooldown_expired(token.id)
+                    # Reload token data after refresh
+                    token = await self.token_manager.db.get_token(token.id)
+
+                # Skip tokens that are in Sora2 cooldown (quota exhausted)
+                if token and token.sora2_cooldown_until and token.sora2_cooldown_until > datetime.now():
+                    continue
+
+                if token:
+                    available_tokens.append(token)
+
+            if not available_tokens:
+                return None
+
+            active_tokens = available_tokens
+
         # If for image generation, filter out locked tokens
         if for_image_generation:
             available_tokens = []

src/services/sora_client.py

@@ -325,3 +325,95 @@ class SoraClient:
                 raise Exception(error_msg)
 
             return True
+
+    async def get_watermark_free_url_custom(self, parse_url: str, parse_token: str, post_id: str) -> str:
+        """Get watermark-free video URL from custom parse server
+
+        Args:
+            parse_url: Custom parse server URL (e.g., http://example.com)
+            parse_token: Access token for custom parse server
+            post_id: Post ID to parse (e.g., s_690c0f574c3881918c3bc5b682a7e9fd)
+
+        Returns:
+            Download link from custom parse server
+
+        Raises:
+            Exception: If parse fails or token is invalid
+        """
+        proxy_url = await self.proxy_manager.get_proxy_url()
+
+        # Construct the share URL
+        share_url = f"https://sora.chatgpt.com/p/{post_id}"
+
+        # Prepare request
+        json_data = {
+            "url": share_url,
+            "token": parse_token
+        }
+
+        kwargs = {
+            "json": json_data,
+            "timeout": 30,
+            "impersonate": "chrome"
+        }
+
+        if proxy_url:
+            kwargs["proxy"] = proxy_url
+
+        try:
+            async with AsyncSession() as session:
+                # Record start time
+                start_time = time.time()
+
+                # Make POST request to custom parse server
+                response = await session.post(f"{parse_url}/get-sora-link", **kwargs)
+
+                # Calculate duration
+                duration_ms = (time.time() - start_time) * 1000
+
+                # Log response
+                debug_logger.log_response(
+                    status_code=response.status_code,
+                    headers=dict(response.headers),
+                    body=response.text if response.text else "No content",
+                    duration_ms=duration_ms
+                )
+
+                # Check status
+                if response.status_code != 200:
+                    error_msg = f"Custom parse failed: {response.status_code} - {response.text}"
+                    debug_logger.log_error(
+                        error_message=error_msg,
+                        status_code=response.status_code,
+                        response_text=response.text
+                    )
+                    raise Exception(error_msg)
+
+                # Parse response
+                result = response.json()
+
+                # Check for error in response
+                if "error" in result:
+                    error_msg = f"Custom parse error: {result['error']}"
+                    debug_logger.log_error(
+                        error_message=error_msg,
+                        status_code=401,
+                        response_text=str(result)
+                    )
+                    raise Exception(error_msg)
+
+                # Extract download link
+                download_link = result.get("download_link")
+                if not download_link:
+                    raise Exception("No download_link in custom parse response")
+
+                debug_logger.log_info(f"Custom parse successful: {download_link}")
+                return download_link
+
+        except Exception as e:
+            debug_logger.log_error(
+                error_message=f"Custom parse request failed: {str(e)}",
+                status_code=500,
+                response_text=str(e)
+            )
+            raise

src/services/token_manager.py

@@ -1,9 +1,11 @@
 """Token management module"""
 import jwt
 import asyncio
+import random
 from datetime import datetime, timedelta
 from typing import Optional, List, Dict, Any
 from curl_cffi.requests import AsyncSession
+from faker import Faker
 from ..core.database import Database
 from ..core.models import Token, TokenStats
 from ..core.config import config
@@ -16,6 +18,7 @@ class TokenManager:
         self.db = db
         self._lock = asyncio.Lock()
         self.proxy_manager = ProxyManager(db)
+        self.fake = Faker()
     
     async def decode_jwt(self, token: str) -> dict:
         """Decode JWT token without verification"""
@@ -24,7 +27,37 @@ class TokenManager:
             return decoded
         except Exception as e:
             raise ValueError(f"Invalid JWT token: {str(e)}")
-    
+
+    def _generate_random_username(self) -> str:
+        """Generate a random username using faker
+
+        Returns:
+            A random username string
+        """
+        # 生成真实姓名
+        first_name = self.fake.first_name()
+        last_name = self.fake.last_name()
+
+        # 去除姓名中的空格和特殊字符，只保留字母
+        first_name_clean = ''.join(c for c in first_name if c.isalpha())
+        last_name_clean = ''.join(c for c in last_name if c.isalpha())
+
+        # 生成1-4位随机数字
+        random_digits = str(random.randint(1, 9999))
+
+        # 随机选择用户名格式
+        format_choice = random.choice([
+            f"{first_name_clean}{last_name_clean}{random_digits}",
+            f"{first_name_clean}.{last_name_clean}{random_digits}",
+            f"{first_name_clean}{random_digits}",
+            f"{last_name_clean}{random_digits}",
+            f"{first_name_clean[0]}{last_name_clean}{random_digits}",
+            f"{first_name_clean}{last_name_clean[0]}{random_digits}"
+        ])
+
+        # 转换为小写
+        return format_choice.lower()
+
     async def get_user_info(self, access_token: str) -> dict:
         """Get user info from Sora API"""
         proxy_url = await self.proxy_manager.get_proxy_url()
@@ -178,6 +211,158 @@ class TokenManager:
                     "invite_code": None
                 }
 
+    async def get_sora2_remaining_count(self, access_token: str) -> dict:
+        """Get Sora2 remaining video count
+
+        Returns:
+            {
+                "remaining_count": 27,
+                "rate_limit_reached": false,
+                "access_resets_in_seconds": 46833
+            }
+        """
+        proxy_url = await self.proxy_manager.get_proxy_url()
+
+        print(f"🔍 开始获取Sora2剩余次数...")
+
+        async with AsyncSession() as session:
+            headers = {
+                "Authorization": f"Bearer {access_token}",
+                "Accept": "application/json"
+            }
+
+            kwargs = {
+                "headers": headers,
+                "timeout": 30,
+                "impersonate": "chrome"  # 自动生成 User-Agent 和浏览器指纹
+            }
+
+            if proxy_url:
+                kwargs["proxy"] = proxy_url
+                print(f"🌐 使用代理: {proxy_url}")
+
+            response = await session.get(
+                "https://sora.chatgpt.com/backend/nf/check",
+                **kwargs
+            )
+
+            print(f"📥 响应状态码: {response.status_code}")
+
+            if response.status_code == 200:
+                data = response.json()
+                print(f"✅ Sora2剩余次数获取成功: {data}")
+
+                rate_limit_info = data.get("rate_limit_and_credit_balance", {})
+                return {
+                    "success": True,
+                    "remaining_count": rate_limit_info.get("estimated_num_videos_remaining", 0),
+                    "rate_limit_reached": rate_limit_info.get("rate_limit_reached", False),
+                    "access_resets_in_seconds": rate_limit_info.get("access_resets_in_seconds", 0)
+                }
+            else:
+                print(f"❌ 获取Sora2剩余次数失败: {response.status_code}")
+                print(f"📄 响应内容: {response.text[:500]}")
+                return {
+                    "success": False,
+                    "remaining_count": 0,
+                    "error": f"Failed to get remaining count: {response.status_code}"
+                }
+
+    async def check_username_available(self, access_token: str, username: str) -> bool:
+        """Check if username is available
+
+        Args:
+            access_token: Access token for authentication
+            username: Username to check
+
+        Returns:
+            True if username is available, False otherwise
+        """
+        proxy_url = await self.proxy_manager.get_proxy_url()
+
+        print(f"🔍 检查用户名是否可用: {username}")
+
+        async with AsyncSession() as session:
+            headers = {
+                "Authorization": f"Bearer {access_token}",
+                "Content-Type": "application/json"
+            }
+
+            kwargs = {
+                "headers": headers,
+                "json": {"username": username},
+                "timeout": 30,
+                "impersonate": "chrome"
+            }
+
+            if proxy_url:
+                kwargs["proxy"] = proxy_url
+                print(f"🌐 使用代理: {proxy_url}")
+
+            response = await session.post(
+                "https://sora.chatgpt.com/backend/project_y/profile/username/check",
+                **kwargs
+            )
+
+            print(f"📥 响应状态码: {response.status_code}")
+
+            if response.status_code == 200:
+                data = response.json()
+                available = data.get("available", False)
+                print(f"✅ 用户名检查结果: available={available}")
+                return available
+            else:
+                print(f"❌ 用户名检查失败: {response.status_code}")
+                print(f"📄 响应内容: {response.text[:500]}")
+                return False
+
+    async def set_username(self, access_token: str, username: str) -> dict:
+        """Set username for the account
+
+        Args:
+            access_token: Access token for authentication
+            username: Username to set
+
+        Returns:
+            User profile information after setting username
+        """
+        proxy_url = await self.proxy_manager.get_proxy_url()
+
+        print(f"🔍 开始设置用户名: {username}")
+
+        async with AsyncSession() as session:
+            headers = {
+                "Authorization": f"Bearer {access_token}",
+                "Content-Type": "application/json"
+            }
+
+            kwargs = {
+                "headers": headers,
+                "json": {"username": username},
+                "timeout": 30,
+                "impersonate": "chrome"
+            }
+
+            if proxy_url:
+                kwargs["proxy"] = proxy_url
+                print(f"🌐 使用代理: {proxy_url}")
+
+            response = await session.post(
+                "https://sora.chatgpt.com/backend/project_y/profile/username/set",
+                **kwargs
+            )
+
+            print(f"📥 响应状态码: {response.status_code}")
+
+            if response.status_code == 200:
+                data = response.json()
+                print(f"✅ 用户名设置成功: {data.get('username')}")
+                return data
+            else:
+                print(f"❌ 用户名设置失败: {response.status_code}")
+                print(f"📄 响应内容: {response.text[:500]}")
+                raise Exception(f"Failed to set username: {response.status_code}")
+
     async def activate_sora2_invite(self, access_token: str, invite_code: str) -> dict:
         """Activate Sora2 with invite code"""
         import uuid
@@ -375,16 +560,63 @@ class TokenManager:
         sora2_invite_code = None
         sora2_redeemed_count = 0
         sora2_total_count = 0
+        sora2_remaining_count = 0
         try:
             sora2_info = await self.get_sora2_invite_code(token_value)
             sora2_supported = sora2_info.get("supported", False)
             sora2_invite_code = sora2_info.get("invite_code")
             sora2_redeemed_count = sora2_info.get("redeemed_count", 0)
             sora2_total_count = sora2_info.get("total_count", 0)
+
+            # If Sora2 is supported, get remaining count
+            if sora2_supported:
+                try:
+                    remaining_info = await self.get_sora2_remaining_count(token_value)
+                    if remaining_info.get("success"):
+                        sora2_remaining_count = remaining_info.get("remaining_count", 0)
+                        print(f"✅ Sora2剩余次数: {sora2_remaining_count}")
+                except Exception as e:
+                    print(f"Failed to get Sora2 remaining count: {e}")
         except Exception as e:
             # If API call fails, Sora2 info will be None
             print(f"Failed to get Sora2 info: {e}")
 
+        # Check and set username if needed
+        try:
+            # Get fresh user info to check username
+            user_info = await self.get_user_info(token_value)
+            username = user_info.get("username")
+
+            # If username is null, need to set one
+            if username is None:
+                print(f"⚠️  检测到用户名为null，需要设置用户名")
+
+                # Generate random username
+                max_attempts = 5
+                for attempt in range(max_attempts):
+                    generated_username = self._generate_random_username()
+                    print(f"🔄 尝试用户名 ({attempt + 1}/{max_attempts}): {generated_username}")
+
+                    # Check if username is available
+                    if await self.check_username_available(token_value, generated_username):
+                        # Set the username
+                        try:
+                            await self.set_username(token_value, generated_username)
+                            print(f"✅ 用户名设置成功: {generated_username}")
+                            break
+                        except Exception as e:
+                            print(f"❌ 用户名设置失败: {e}")
+                            if attempt == max_attempts - 1:
+                                print(f"⚠️  达到最大尝试次数，跳过用户名设置")
+                    else:
+                        print(f"⚠️  用户名 {generated_username} 已被占用，尝试下一个")
+                        if attempt == max_attempts - 1:
+                            print(f"⚠️  达到最大尝试次数，跳过用户名设置")
+            else:
+                print(f"✅ 用户名已设置: {username}")
+        except Exception as e:
+            print(f"⚠️  用户名检查/设置过程中出错: {e}")
+
         # Create token object
         token = Token(
             token=token_value,
@@ -401,7 +633,8 @@ class TokenManager:
             sora2_supported=sora2_supported,
             sora2_invite_code=sora2_invite_code,
             sora2_redeemed_count=sora2_redeemed_count,
-            sora2_total_count=sora2_total_count
+            sora2_total_count=sora2_total_count,
+            sora2_remaining_count=sora2_remaining_count
         )
 
         # Save to database
@@ -523,6 +756,16 @@ class TokenManager:
             sora2_invite_code = sora2_info.get("invite_code")
             sora2_redeemed_count = sora2_info.get("redeemed_count", 0)
             sora2_total_count = sora2_info.get("total_count", 0)
+            sora2_remaining_count = 0
+
+            # If Sora2 is supported, get remaining count
+            if sora2_supported:
+                try:
+                    remaining_info = await self.get_sora2_remaining_count(token_data.token)
+                    if remaining_info.get("success"):
+                        sora2_remaining_count = remaining_info.get("remaining_count", 0)
+                except Exception as e:
+                    print(f"Failed to get Sora2 remaining count: {e}")
 
             # Update token Sora2 info in database
             await self.db.update_token_sora2(
@@ -530,7 +773,8 @@ class TokenManager:
                 supported=sora2_supported,
                 invite_code=sora2_invite_code,
                 redeemed_count=sora2_redeemed_count,
-                total_count=sora2_total_count
+                total_count=sora2_total_count,
+                remaining_count=sora2_remaining_count
             )
 
             return {
@@ -541,7 +785,8 @@ class TokenManager:
                 "sora2_supported": sora2_supported,
                 "sora2_invite_code": sora2_invite_code,
                 "sora2_redeemed_count": sora2_redeemed_count,
-                "sora2_total_count": sora2_total_count
+                "sora2_total_count": sora2_total_count,
+                "sora2_remaining_count": sora2_remaining_count
             }
         except Exception as e:
             return {
@@ -569,16 +814,51 @@ class TokenManager:
         if stats and stats.error_count >= admin_config.error_ban_threshold:
             await self.db.update_token_status(token_id, False)
     
-    async def record_success(self, token_id: int):
+    async def record_success(self, token_id: int, is_video: bool = False):
         """Record successful request (reset error count)"""
         await self.db.reset_error_count(token_id)
+
+        # Update Sora2 remaining count after video generation
+        if is_video:
+            try:
+                token_data = await self.db.get_token(token_id)
+                if token_data and token_data.sora2_supported:
+                    remaining_info = await self.get_sora2_remaining_count(token_data.token)
+                    if remaining_info.get("success"):
+                        remaining_count = remaining_info.get("remaining_count", 0)
+                        await self.db.update_token_sora2_remaining(token_id, remaining_count)
+                        print(f"✅ 更新Token {token_id} 的Sora2剩余次数: {remaining_count}")
+
+                        # If remaining count is 0, set cooldown
+                        if remaining_count == 0:
+                            reset_seconds = remaining_info.get("access_resets_in_seconds", 0)
+                            if reset_seconds > 0:
+                                cooldown_until = datetime.now() + timedelta(seconds=reset_seconds)
+                                await self.db.update_token_sora2_cooldown(token_id, cooldown_until)
+                                print(f"⏱️ Token {token_id} 剩余次数为0，设置冷却时间至: {cooldown_until}")
+            except Exception as e:
+                print(f"Failed to update Sora2 remaining count: {e}")
     
-    async def check_and_apply_cooldown(self, token_id: int):
-        """Check if token should be cooled down"""
-        stats = await self.db.get_token_stats(token_id)
-        admin_config = await self.db.get_admin_config()
-        
-        if stats and stats.video_count >= admin_config.video_cooldown_threshold:
-            # Apply 12 hour cooldown
-            cooled_until = datetime.now() + timedelta(hours=12)
-            await self.db.update_token_cooldown(token_id, cooled_until)
+    async def refresh_sora2_remaining_if_cooldown_expired(self, token_id: int):
+        """Refresh Sora2 remaining count if cooldown has expired"""
+        try:
+            token_data = await self.db.get_token(token_id)
+            if not token_data or not token_data.sora2_supported:
+                return
+
+            # Check if Sora2 cooldown has expired
+            if token_data.sora2_cooldown_until and token_data.sora2_cooldown_until <= datetime.now():
+                print(f"🔄 Token {token_id} Sora2冷却已过期，正在刷新剩余次数...")
+
+                try:
+                    remaining_info = await self.get_sora2_remaining_count(token_data.token)
+                    if remaining_info.get("success"):
+                        remaining_count = remaining_info.get("remaining_count", 0)
+                        await self.db.update_token_sora2_remaining(token_id, remaining_count)
+                        # Clear cooldown
+                        await self.db.update_token_sora2_cooldown(token_id, None)
+                        print(f"✅ Token {token_id} Sora2剩余次数已刷新: {remaining_count}")
+                except Exception as e:
+                    print(f"Failed to refresh Sora2 remaining count: {e}")
+        except Exception as e:
+            print(f"Error in refresh_sora2_remaining_if_cooldown_expired: {e}")