1

video-gen-api/app/models/chat_generation_task.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from sqlalchemy import DateTime, Float, ForeignKey, Index, Integer, String, Text, text
+from sqlalchemy import CheckConstraint, DateTime, Float, ForeignKey, Index, Integer, String, Text, text
 from sqlalchemy.orm import Mapped, mapped_column
 
 from app.models.base import Base, TimestampMixin, SoftDeleteMixin
@@ -26,6 +26,19 @@ class ChatGenerationTask(Base, TimestampMixin, SoftDeleteMixin):
             unique=True,
             postgresql_where=text("deleted_at IS NULL AND idempotency_key IS NOT NULL"),
         ),
+        # AI 创作顶层任务在 chatapi_async/chatapi_main 之间切换时，
+        # 同一个前端幂等键也只能创建一组任务。
+        Index(
+            "uq_chat_generation_tasks_user_chat_idempotency",
+            "user_id",
+            "idempotency_key",
+            unique=True,
+            postgresql_where=text(
+                "deleted_at IS NULL "
+                "AND idempotency_key IS NOT NULL "
+                "AND generation_mode IN ('chatapi_async', 'chatapi_main')"
+            ),
+        ),
         # 视频 24 小时降频轮询调度使用。
         Index(
             "idx_chat_generation_tasks_next_poll_at",
@@ -37,6 +50,17 @@ class ChatGenerationTask(Base, TimestampMixin, SoftDeleteMixin):
                 "AND next_poll_at IS NOT NULL"
             ),
         ),
+        Index(
+            "uq_chat_generation_tasks_parent_index",
+            "parent_task_id",
+            "generation_index",
+            unique=True,
+            postgresql_where=text("parent_task_id IS NOT NULL AND generation_index IS NOT NULL"),
+        ),
+        Index("idx_chat_generation_tasks_parent", "parent_task_id"),
+        Index("idx_chat_generation_tasks_user_mode_created", "user_id", "generation_mode", "created_at"),
+        CheckConstraint("generation_count BETWEEN 1 AND 5", name="ck_chat_generation_tasks_generation_count"),
+        CheckConstraint("generation_index IS NULL OR generation_index > 0", name="ck_chat_generation_tasks_generation_index"),
     )
 
 
@@ -59,6 +83,17 @@ class ChatGenerationTask(Base, TimestampMixin, SoftDeleteMixin):
     status: Mapped[str] = mapped_column(String(32), default="generating", index=True)
     pipeline_stage: Mapped[str | None] = mapped_column(String(32), nullable=True, index=True)
     generation_mode: Mapped[str] = mapped_column(String(32), default="chatapi_async", index=True)
+    parent_task_id: Mapped[str | None] = mapped_column(
+        String(32), ForeignKey("chat_generation_tasks.id", ondelete="RESTRICT"), nullable=True
+    )
+    generation_count: Mapped[int] = mapped_column(Integer, default=1, server_default="1", nullable=False)
+    generation_index: Mapped[int | None] = mapped_column(Integer, nullable=True)
+
+    # 图片主任务同步调用供应商时的分布式执行租约。
+    # 防止重复 Celery 消息或恢复任务同时触发多次组图请求。
+    provider_create_claim_token: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
+    provider_create_lease_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
+    provider_create_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
 
     media_references: Mapped[str | None] = mapped_column(Text, nullable=True)
     provider_task_id: Mapped[str | None] = mapped_column(String(128), nullable=True, index=True)