拆镜复刻、爆款开头复刻容灾优化

video-gen-api/app/tasks/hot_opening_replicate_tasks.py

@@ -1,45 +1,95 @@
 from __future__ import annotations
 
+from typing import Any
+
+from app.enums.hot_opening_replicate import HotOpeningStepCodeEnum, ModuleCodeEnum
 from app.models.base import async_session
 from app.services.hot_opening_replicate_service import run_image_prompt_optimize, run_video_prompt_optimize
+from app.services.module_async_recovery_service import (
+    OBJECT_MODULE_STEP,
+    TASK_HOT_IMAGE_PROMPT,
+    TASK_HOT_VIDEO_PROMPT,
+    acquire_object_lock,
+    cleanup_active_if_terminal,
+    mark_active_started,
+    release_object_lock,
+    register_module_step_task,
+)
 from app.tasks.async_runner import run_async
 from app.tasks.celery_app import celery_app
 
+MODULE = ModuleCodeEnum.HOT_OPENING_REPLICATE.value
+
 
 async def _run_image_prompt(project_id: str, step_id: str | None = None):
-    async with async_session() as db:
-        await run_image_prompt_optimize(db, project_id=project_id, step_id=step_id)
-        await db.commit()
+    lock_token: str | None = None
+    if step_id:
+        lock_token = await acquire_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id)
+        if not lock_token:
+            # 同一个 step 的重复消息正在被其它 worker 处理：当前消息直接忽略，active registry 保留等待真正任务完成或 lease 过期。
+            return None
+        await register_module_step_task(
+            module=MODULE,
+            project_id=project_id,
+            step_id=step_id,
+            step_code=HotOpeningStepCodeEnum.IMAGE_PROMPT_OPTIMIZE.value,
+            task_name=TASK_HOT_IMAGE_PROMPT,
+        )
+        await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
+    try:
+        async with async_session() as db:
+            result = await run_image_prompt_optimize(db, project_id=project_id, step_id=step_id)
+            await db.commit()
+            if step_id:
+                await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
+            return result
+    finally:
+        if step_id:
+            await release_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id, token=lock_token)
 
 
 async def _run_video_prompt(project_id: str, step_id: str | None = None):
-    async with async_session() as db:
-        await run_video_prompt_optimize(db, project_id=project_id, step_id=step_id)
-        await db.commit()
+    lock_token: str | None = None
+    if step_id:
+        lock_token = await acquire_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id)
+        if not lock_token:
+            return None
+        await register_module_step_task(
+            module=MODULE,
+            project_id=project_id,
+            step_id=step_id,
+            step_code=HotOpeningStepCodeEnum.VIDEO_PROMPT_OPTIMIZE.value,
+            task_name=TASK_HOT_VIDEO_PROMPT,
+        )
+        await mark_active_started(object_type=OBJECT_MODULE_STEP, object_id=step_id)
+    try:
+        async with async_session() as db:
+            result = await run_video_prompt_optimize(db, project_id=project_id, step_id=step_id)
+            await db.commit()
+            if step_id:
+                await cleanup_active_if_terminal(db, object_type=OBJECT_MODULE_STEP, object_id=step_id)
+            return result
+    finally:
+        if step_id:
+            await release_object_lock(object_type=OBJECT_MODULE_STEP, object_id=step_id, token=lock_token)
 
 
 if celery_app:
     @celery_app.task(name="hot_opening.start_image_prompt_optimize", bind=True, max_retries=3, default_retry_delay=30)
     def start_image_prompt_optimize(self, project_id: str, step_id: str | None = None):
-        """手动触发后的图片 AI 提词任务。
-
-        该任务路由到现有 gen_chatapi_create 队列，不需要新增 hot_opening worker。
-        """
+        """手动触发后的图片 AI 提词任务。"""
         return run_async(_run_image_prompt(project_id, step_id))
 
     @celery_app.task(name="hot_opening.start_video_prompt_optimize", bind=True, max_retries=3, default_retry_delay=30)
     def start_video_prompt_optimize(self, project_id: str, step_id: str | None = None):
-        """手动触发后的视频 AI 提词任务。
-
-        该任务路由到现有 gen_chatapi_create 队列，不需要新增 hot_opening worker。
-        """
+        """手动触发后的视频 AI 提词任务。"""
         return run_async(_run_video_prompt(project_id, step_id))
 else:
     class _DisabledTask:
-        def delay(self, *args, **kwargs):
+        def delay(self, *args: Any, **kwargs: Any):
             raise RuntimeError("Celery is disabled")
 
-        def apply_async(self, *args, **kwargs):
+        def apply_async(self, *args: Any, **kwargs: Any):
             raise RuntimeError("Celery is disabled")
 
     start_image_prompt_optimize = _DisabledTask()