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

src/api/admin.py

@@ -79,7 +79,6 @@ class UpdateTokenRequest(BaseModel):
     remark: Optional[str] = None
 
 class UpdateAdminConfigRequest(BaseModel):
-    video_cooldown_threshold: int
     error_ban_threshold: int
 
 class UpdateProxyConfigRequest(BaseModel):
@@ -109,6 +108,9 @@ class UpdateGenerationTimeoutRequest(BaseModel):
 
 class UpdateWatermarkFreeConfigRequest(BaseModel):
     watermark_free_enabled: bool
+    parse_method: Optional[str] = "third_party"  # "third_party" or "custom"
+    custom_parse_url: Optional[str] = None
+    custom_parse_token: Optional[str] = None
 
 class UpdateVideoLengthConfigRequest(BaseModel):
     default_length: str  # "10s" or "15s"
@@ -139,7 +141,7 @@ async def get_tokens(token: str = Depends(verify_admin_token)) -> List[dict]:
     """Get all tokens with statistics"""
     tokens = await token_manager.get_all_tokens()
     result = []
-    
+
     for token in tokens:
         stats = await db.get_token_stats(token.id)
         result.append({
@@ -167,9 +169,11 @@ async def get_tokens(token: str = Depends(verify_admin_token)) -> List[dict]:
             "sora2_supported": token.sora2_supported,
             "sora2_invite_code": token.sora2_invite_code,
             "sora2_redeemed_count": token.sora2_redeemed_count,
-            "sora2_total_count": token.sora2_total_count
+            "sora2_total_count": token.sora2_total_count,
+            "sora2_remaining_count": token.sora2_remaining_count,
+            "sora2_cooldown_until": token.sora2_cooldown_until.isoformat() if token.sora2_cooldown_until else None
         })
-    
+
     return result
 
 @router.post("/api/tokens")
@@ -275,7 +279,8 @@ async def test_token(token_id: int, token: str = Depends(verify_admin_token)):
                 "sora2_supported": result.get("sora2_supported"),
                 "sora2_invite_code": result.get("sora2_invite_code"),
                 "sora2_redeemed_count": result.get("sora2_redeemed_count"),
-                "sora2_total_count": result.get("sora2_total_count")
+                "sora2_total_count": result.get("sora2_total_count"),
+                "sora2_remaining_count": result.get("sora2_remaining_count")
             })
 
         return response
@@ -316,7 +321,6 @@ async def get_admin_config(token: str = Depends(verify_admin_token)) -> dict:
     """Get admin configuration"""
     admin_config = await db.get_admin_config()
     return {
-        "video_cooldown_threshold": admin_config.video_cooldown_threshold,
         "error_ban_threshold": admin_config.error_ban_threshold,
         "api_key": config.api_key,
         "admin_username": config.admin_username,
@@ -331,7 +335,6 @@ async def update_admin_config(
     """Update admin configuration"""
     try:
         admin_config = AdminConfig(
-            video_cooldown_threshold=request.video_cooldown_threshold,
             error_ban_threshold=request.error_ban_threshold
         )
         await db.update_admin_config(admin_config)
@@ -480,9 +483,12 @@ async def update_proxy_config(
 @router.get("/api/watermark-free/config")
 async def get_watermark_free_config(token: str = Depends(verify_admin_token)) -> dict:
     """Get watermark-free mode configuration"""
-    config = await db.get_watermark_free_config()
+    config_obj = await db.get_watermark_free_config()
     return {
-        "watermark_free_enabled": config.watermark_free_enabled
+        "watermark_free_enabled": config_obj.watermark_free_enabled,
+        "parse_method": config_obj.parse_method,
+        "custom_parse_url": config_obj.custom_parse_url,
+        "custom_parse_token": config_obj.custom_parse_token
     }
 
 @router.post("/api/watermark-free/config")
@@ -492,7 +498,12 @@ async def update_watermark_free_config(
 ):
     """Update watermark-free mode configuration"""
     try:
-        await db.update_watermark_free_config(request.watermark_free_enabled)
+        await db.update_watermark_free_config(
+            request.watermark_free_enabled,
+            request.parse_method,
+            request.custom_parse_url,
+            request.custom_parse_token
+        )
 
         # Update in-memory config
         from ..core.config import config
@@ -549,13 +560,23 @@ async def activate_sora2(
             # Get new invite code after activation
             sora2_info = await token_manager.get_sora2_invite_code(token_obj.token)
 
+            # Get remaining count
+            sora2_remaining_count = 0
+            try:
+                remaining_info = await token_manager.get_sora2_remaining_count(token_obj.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 database
             await db.update_token_sora2(
                 token_id,
                 supported=True,
                 invite_code=sora2_info.get("invite_code"),
                 redeemed_count=sora2_info.get("redeemed_count", 0),
-                total_count=sora2_info.get("total_count", 0)
+                total_count=sora2_info.get("total_count", 0),
+                remaining_count=sora2_remaining_count
             )
 
             return {
@@ -564,7 +585,8 @@ async def activate_sora2(
                 "already_accepted": result.get("already_accepted", False),
                 "invite_code": sora2_info.get("invite_code"),
                 "redeemed_count": sora2_info.get("redeemed_count", 0),
-                "total_count": sora2_info.get("total_count", 0)
+                "total_count": sora2_info.get("total_count", 0),
+                "sora2_remaining_count": sora2_remaining_count
             }
         else:
             return {

src/core/config.py

@@ -145,13 +145,28 @@ class Config:
     @property
     def watermark_free_enabled(self) -> bool:
         """Get watermark-free mode enabled status"""
-        return self._config.get("watermark_free", {}).get("enabled", False)
+        return self._config.get("watermark_free", {}).get("watermark_free_enabled", False)
 
     def set_watermark_free_enabled(self, enabled: bool):
         """Set watermark-free mode enabled/disabled"""
         if "watermark_free" not in self._config:
             self._config["watermark_free"] = {}
-        self._config["watermark_free"]["enabled"] = enabled
+        self._config["watermark_free"]["watermark_free_enabled"] = enabled
+
+    @property
+    def watermark_free_parse_method(self) -> str:
+        """Get watermark-free parse method"""
+        return self._config.get("watermark_free", {}).get("parse_method", "third_party")
+
+    @property
+    def watermark_free_custom_url(self) -> str:
+        """Get custom parse server URL"""
+        return self._config.get("watermark_free", {}).get("custom_parse_url", "")
+
+    @property
+    def watermark_free_custom_token(self) -> str:
+        """Get custom parse server access token"""
+        return self._config.get("watermark_free", {}).get("custom_parse_token", "")
 
 # Global config instance
 config = Config()

src/core/database.py

@@ -47,7 +47,9 @@ class Database:
                     sora2_supported BOOLEAN,
                     sora2_invite_code TEXT,
                     sora2_redeemed_count INTEGER DEFAULT 0,
-                    sora2_total_count INTEGER DEFAULT 0
+                    sora2_total_count INTEGER DEFAULT 0,
+                    sora2_remaining_count INTEGER DEFAULT 0,
+                    sora2_cooldown_until TIMESTAMP
                 )
             """)
 
@@ -71,7 +73,33 @@ class Database:
                 await db.execute("ALTER TABLE tokens ADD COLUMN sora2_total_count INTEGER DEFAULT 0")
             except:
                 pass  # Column already exists
-            
+
+            try:
+                await db.execute("ALTER TABLE tokens ADD COLUMN sora2_remaining_count INTEGER DEFAULT 0")
+            except:
+                pass  # Column already exists
+
+            try:
+                await db.execute("ALTER TABLE tokens ADD COLUMN sora2_cooldown_until TIMESTAMP")
+            except:
+                pass  # Column already exists
+
+            # Migrate watermark_free_config table - add new columns
+            try:
+                await db.execute("ALTER TABLE watermark_free_config ADD COLUMN parse_method TEXT DEFAULT 'third_party'")
+            except:
+                pass  # Column already exists
+
+            try:
+                await db.execute("ALTER TABLE watermark_free_config ADD COLUMN custom_parse_url TEXT")
+            except:
+                pass  # Column already exists
+
+            try:
+                await db.execute("ALTER TABLE watermark_free_config ADD COLUMN custom_parse_token TEXT")
+            except:
+                pass  # Column already exists
+
             # Token stats table
             await db.execute("""
                 CREATE TABLE IF NOT EXISTS token_stats (
@@ -122,7 +150,6 @@ class Database:
             await db.execute("""
                 CREATE TABLE IF NOT EXISTS admin_config (
                     id INTEGER PRIMARY KEY DEFAULT 1,
-                    video_cooldown_threshold INTEGER DEFAULT 30,
                     error_ban_threshold INTEGER DEFAULT 3,
                     updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                 )
@@ -144,6 +171,9 @@ class Database:
                 CREATE TABLE IF NOT EXISTS watermark_free_config (
                     id INTEGER PRIMARY KEY DEFAULT 1,
                     watermark_free_enabled BOOLEAN DEFAULT 0,
+                    parse_method TEXT DEFAULT 'third_party',
+                    custom_parse_url TEXT,
+                    custom_parse_token TEXT,
                     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                     updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                 )
@@ -167,8 +197,8 @@ class Database:
             
             # Insert default admin config
             await db.execute("""
-                INSERT OR IGNORE INTO admin_config (id, video_cooldown_threshold, error_ban_threshold)
-                VALUES (1, 30, 3)
+                INSERT OR IGNORE INTO admin_config (id, error_ban_threshold)
+                VALUES (1, 3)
             """)
             
             # Insert default proxy config
@@ -179,8 +209,8 @@ class Database:
 
             # Insert default watermark-free config
             await db.execute("""
-                INSERT OR IGNORE INTO watermark_free_config (id, watermark_free_enabled)
-                VALUES (1, 0)
+                INSERT OR IGNORE INTO watermark_free_config (id, watermark_free_enabled, parse_method, custom_parse_url, custom_parse_token)
+                VALUES (1, 0, 'third_party', NULL, NULL)
             """)
 
             # Insert default video length config
@@ -196,14 +226,13 @@ class Database:
         async with aiosqlite.connect(self.db_path) as db:
             # Initialize admin config
             admin_config = config_dict.get("admin", {})
-            video_cooldown_threshold = admin_config.get("video_cooldown_threshold", 30)
             error_ban_threshold = admin_config.get("error_ban_threshold", 3)
 
             await db.execute("""
                 UPDATE admin_config
-                SET video_cooldown_threshold = ?, error_ban_threshold = ?, updated_at = CURRENT_TIMESTAMP
+                SET error_ban_threshold = ?, updated_at = CURRENT_TIMESTAMP
                 WHERE id = 1
-            """, (video_cooldown_threshold, error_ban_threshold))
+            """, (error_ban_threshold,))
 
             # Initialize proxy config
             proxy_config = config_dict.get("proxy", {})
@@ -221,12 +250,20 @@ class Database:
             # Initialize watermark-free config
             watermark_config = config_dict.get("watermark_free", {})
             watermark_free_enabled = watermark_config.get("watermark_free_enabled", False)
+            parse_method = watermark_config.get("parse_method", "third_party")
+            custom_parse_url = watermark_config.get("custom_parse_url", "")
+            custom_parse_token = watermark_config.get("custom_parse_token", "")
+
+            # Convert empty strings to None
+            custom_parse_url = custom_parse_url if custom_parse_url else None
+            custom_parse_token = custom_parse_token if custom_parse_token else None
 
             await db.execute("""
                 UPDATE watermark_free_config
-                SET watermark_free_enabled = ?, updated_at = CURRENT_TIMESTAMP
+                SET watermark_free_enabled = ?, parse_method = ?, custom_parse_url = ?,
+                    custom_parse_token = ?, updated_at = CURRENT_TIMESTAMP
                 WHERE id = 1
-            """, (watermark_free_enabled,))
+            """, (watermark_free_enabled, parse_method, custom_parse_url, custom_parse_token))
 
             # Initialize video length config
             video_length_config = config_dict.get("video_length", {})
@@ -249,13 +286,14 @@ class Database:
             cursor = await db.execute("""
                 INSERT INTO tokens (token, email, username, name, st, rt, remark, expiry_time, is_active,
                                    plan_type, plan_title, subscription_end, sora2_supported, sora2_invite_code,
-                                   sora2_redeemed_count, sora2_total_count)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                                   sora2_redeemed_count, sora2_total_count, sora2_remaining_count, sora2_cooldown_until)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
             """, (token.token, token.email, "", token.name, token.st, token.rt,
                   token.remark, token.expiry_time, token.is_active,
                   token.plan_type, token.plan_title, token.subscription_end,
                   token.sora2_supported, token.sora2_invite_code,
-                  token.sora2_redeemed_count, token.sora2_total_count))
+                  token.sora2_redeemed_count, token.sora2_total_count,
+                  token.sora2_remaining_count, token.sora2_cooldown_until))
             await db.commit()
             token_id = cursor.lastrowid
 
@@ -328,14 +366,30 @@ class Database:
             await db.commit()
     
     async def update_token_sora2(self, token_id: int, supported: bool, invite_code: Optional[str] = None,
-                                redeemed_count: int = 0, total_count: int = 0):
+                                redeemed_count: int = 0, total_count: int = 0, remaining_count: int = 0):
         """Update token Sora2 support info"""
         async with aiosqlite.connect(self.db_path) as db:
             await db.execute("""
                 UPDATE tokens
-                SET sora2_supported = ?, sora2_invite_code = ?, sora2_redeemed_count = ?, sora2_total_count = ?
+                SET sora2_supported = ?, sora2_invite_code = ?, sora2_redeemed_count = ?, sora2_total_count = ?, sora2_remaining_count = ?
                 WHERE id = ?
-            """, (supported, invite_code, redeemed_count, total_count, token_id))
+            """, (supported, invite_code, redeemed_count, total_count, remaining_count, token_id))
+            await db.commit()
+
+    async def update_token_sora2_remaining(self, token_id: int, remaining_count: int):
+        """Update token Sora2 remaining count"""
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute("""
+                UPDATE tokens SET sora2_remaining_count = ? WHERE id = ?
+            """, (remaining_count, token_id))
+            await db.commit()
+
+    async def update_token_sora2_cooldown(self, token_id: int, cooldown_until: Optional[datetime]):
+        """Update token Sora2 cooldown time"""
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute("""
+                UPDATE tokens SET sora2_cooldown_until = ? WHERE id = ?
+            """, (cooldown_until, token_id))
             await db.commit()
 
     async def update_token_cooldown(self, token_id: int, cooled_until: datetime):
@@ -533,10 +587,10 @@ class Database:
         """Update admin configuration"""
         async with aiosqlite.connect(self.db_path) as db:
             await db.execute("""
-                UPDATE admin_config 
-                SET video_cooldown_threshold = ?, error_ban_threshold = ?, updated_at = CURRENT_TIMESTAMP
+                UPDATE admin_config
+                SET error_ban_threshold = ?, updated_at = CURRENT_TIMESTAMP
                 WHERE id = 1
-            """, (config.video_cooldown_threshold, config.error_ban_threshold))
+            """, (config.error_ban_threshold,))
             await db.commit()
     
     # Proxy config operations
@@ -571,14 +625,25 @@ class Database:
                 return WatermarkFreeConfig(**dict(row))
             return WatermarkFreeConfig()
 
-    async def update_watermark_free_config(self, enabled: bool):
+    async def update_watermark_free_config(self, enabled: bool, parse_method: str = None,
+                                          custom_parse_url: str = None, custom_parse_token: str = None):
         """Update watermark-free configuration"""
         async with aiosqlite.connect(self.db_path) as db:
-            await db.execute("""
-                UPDATE watermark_free_config
-                SET watermark_free_enabled = ?, updated_at = CURRENT_TIMESTAMP
-                WHERE id = 1
-            """, (enabled,))
+            if parse_method is None and custom_parse_url is None and custom_parse_token is None:
+                # Only update enabled status
+                await db.execute("""
+                    UPDATE watermark_free_config
+                    SET watermark_free_enabled = ?, updated_at = CURRENT_TIMESTAMP
+                    WHERE id = 1
+                """, (enabled,))
+            else:
+                # Update all fields
+                await db.execute("""
+                    UPDATE watermark_free_config
+                    SET watermark_free_enabled = ?, parse_method = ?, custom_parse_url = ?,
+                        custom_parse_token = ?, updated_at = CURRENT_TIMESTAMP
+                    WHERE id = 1
+                """, (enabled, parse_method or "third_party", custom_parse_url, custom_parse_token))
             await db.commit()
 
     # Video length config operations

src/core/models.py

@@ -27,6 +27,9 @@ class Token(BaseModel):
     sora2_invite_code: Optional[str] = None  # Sora2邀请码
     sora2_redeemed_count: int = 0  # Sora2已用次数
     sora2_total_count: int = 0  # Sora2总次数
+    # Sora2 剩余次数
+    sora2_remaining_count: int = 0  # Sora2剩余可用次数
+    sora2_cooldown_until: Optional[datetime] = None  # Sora2冷却时间
 
 class TokenStats(BaseModel):
     """Token statistics"""
@@ -65,7 +68,6 @@ class RequestLog(BaseModel):
 class AdminConfig(BaseModel):
     """Admin configuration"""
     id: int = 1
-    video_cooldown_threshold: int = 30
     error_ban_threshold: int = 3
     updated_at: Optional[datetime] = None
 
@@ -81,6 +83,9 @@ class WatermarkFreeConfig(BaseModel):
     """Watermark-free mode configuration"""
     id: int = 1
     watermark_free_enabled: bool = False
+    parse_method: str = "third_party"  # "third_party" or "custom"
+    custom_parse_url: Optional[str] = None  # Custom parse server URL
+    custom_parse_token: Optional[str] = None  # Custom parse server access token
     created_at: Optional[datetime] = None
     updated_at: Optional[datetime] = None
 

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}")