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

from app.tasks.async_runner import run_async
from datetime import datetime, timezone, timedelta
import uuid

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.models.base import async_session
from app.models.chat_generation_task import ChatGenerationTask
from app.services.celery_download_recovery_service import (
    build_download_active_payload,
    ensure_aware_utc,
    remove_download_active,
    upsert_download_active,
)
from app.services.error_codes import extract_error_message
from app.services.generation_download_service import download_generation_result
from app.services.generation_log_service import log_task_event
from app.services.generation_refund_service import mark_chat_generation_task_failed_and_refund_once
from app.services.resource_accounting_service import record_chat_task_generated_resource
from app.tasks.celery_app import celery_app

DOWNLOAD_QUEUE = "gen_result_download"
DOWNLOAD_STAGE_QUEUED = "download_queued"
DOWNLOAD_STAGE_DOWNLOADING = "downloading"
DOWNLOAD_STAGE_RETRY_WAITING = "retry_waiting"
DOWNLOAD_STAGE_DONE = "done"
DOWNLOAD_STAGE_FAILED = "download_failed"


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _task_created_date_dir(task: ChatGenerationTask) -> str:
    created_at = ensure_aware_utc(getattr(task, "created_at", None)) or _now()
    return created_at.strftime("%Y/%m/%d")


def _queue_timeout_at(now: datetime | None = None) -> datetime:
    now = now or _now()
    return now + timedelta(seconds=int(settings.DOWNLOAD_TASK_QUEUE_TIMEOUT_SECONDS or 300))


def _lease_until(now: datetime | None = None) -> datetime:
    now = now or _now()
    return now + timedelta(seconds=int(settings.DOWNLOAD_TASK_LEASE_SECONDS or 600))


def _retry_at(attempt: int, now: datetime | None = None) -> datetime:
    now = now or _now()
    base = int(settings.DOWNLOAD_TASK_RETRY_BACKOFF_SECONDS or 30)
    return now + timedelta(seconds=max(1, base * max(1, attempt)))


def _is_expired(value: datetime | None, now: datetime | None = None) -> bool:
    value = ensure_aware_utc(value)
    if value is None:
        return True
    return value <= (now or _now())


def _is_already_completed(task: ChatGenerationTask) -> bool:
    if task.status == "completed" or task.pipeline_stage == DOWNLOAD_STAGE_DONE:
        if task.gen_type == "image" and task.image_url:
            return True
        if task.gen_type == "video" and task.video_url:
            return True
    return False


def _build_celery_task_id(task_id: str, attempt: int | None = None, reason: str | None = None) -> str:
    safe_reason = (reason or "download").replace(" ", "_")[:32]
    return f"download:{task_id}:{int(attempt or 0)}:{safe_reason}:{uuid.uuid4().hex[:12]}"


async def _register_active_from_task(
    task: ChatGenerationTask,
    *,
    check_at: datetime,
    priority: int,
    reason: str | None = None,
) -> None:
    payload = build_download_active_payload(
        record_id=task.id,
        celery_task_id=task.download_celery_task_id,
        stage=task.pipeline_stage or "",
        attempt=task.download_attempt_count or 0,
        queue=DOWNLOAD_QUEUE,
        priority=priority,
        enqueue_at=task.download_enqueued_at,
        started_at=task.download_started_at,
        lease_until=task.download_lease_until,
        next_retry_at=task.download_next_retry_at,
        check_at=check_at,
        reason=reason,
    )
    await upsert_download_active(record_id=task.id, payload=payload, check_at=check_at)


async def enqueue_download_task(
    db: AsyncSession,
    task: ChatGenerationTask,
    *,
    recover: bool = False,
    reason: str | None = None,
    countdown: int | None = None,
) -> str | None:
    """统一投递图片/视频下载任务，并同步 DB + Redis active 注册表。"""
    if not task or task.generation_mode != "chatapi_async":
        return None
    if task.status != "generating":
        return None
    if not task.remote_result_url:
        return None

    now = _now()
    priority = settings.DOWNLOAD_TASK_PRIORITY_RECOVER if recover else settings.DOWNLOAD_TASK_PRIORITY_NORMAL
    celery_task_id = _build_celery_task_id(
        task.id,
        attempt=task.download_attempt_count or task.retry_count or 0,
        reason=reason or ("recover" if recover else "normal"),
    )

    task.pipeline_stage = DOWNLOAD_STAGE_QUEUED
    task.download_celery_task_id = celery_task_id
    task.download_enqueued_at = now
    task.download_next_retry_at = None
    if not task.download_storage_date_dir:
        task.download_storage_date_dir = _task_created_date_dir(task)

    await db.commit()

    check_at = _queue_timeout_at(now)
    await _register_active_from_task(task, check_at=check_at, priority=priority, reason=reason)

    if celery_app:
        download_generation_result_task.apply_async(
            args=[task.id],
            queue=DOWNLOAD_QUEUE,
            priority=priority,
            countdown=countdown,
            task_id=celery_task_id,
        )
    return celery_task_id


async def _reload_task(db: AsyncSession, task_id: str) -> ChatGenerationTask | None:
    result = await db.execute(
        select(ChatGenerationTask).where(
            ChatGenerationTask.id == task_id,
            ChatGenerationTask.deleted_at.is_(None),
        ).with_for_update().limit(1)
    )
    return result.scalar_one_or_none()


async def _claim_download_lease(db: AsyncSession, task: ChatGenerationTask) -> bool:
    now = _now()

    if not task or task.generation_mode != "chatapi_async":
        return False
    if task.status != "generating":
        return False
    if _is_already_completed(task):
        return False
    if not task.remote_result_url:
        return False

    stage = task.pipeline_stage

    if stage == DOWNLOAD_STAGE_DOWNLOADING:
        if not _is_expired(task.download_lease_until, now):
            return False
        await log_task_event(
            task,
            event_type="DOWNLOAD_STUCK_RECOVER",
            message=f"downloading lease 已过期，重新抢占下载。lease_until={task.download_lease_until}",
        )
    elif stage == DOWNLOAD_STAGE_RETRY_WAITING:
        if not _is_expired(task.download_next_retry_at, now):
            return False
    elif stage in (DOWNLOAD_STAGE_QUEUED, "result_ready"):
        pass
    else:
        return False

    old_stage = stage
    task.pipeline_stage = DOWNLOAD_STAGE_DOWNLOADING
    task.download_started_at = now
    task.download_lease_until = _lease_until(now)
    task.download_next_retry_at = None
    task.download_attempt_count = int(task.download_attempt_count or 0) + 1
    task.retry_count = task.download_attempt_count
    if not task.download_storage_date_dir:
        task.download_storage_date_dir = _task_created_date_dir(task)

    await db.commit()

    await _register_active_from_task(
        task,
        check_at=task.download_lease_until,
        priority=settings.DOWNLOAD_TASK_PRIORITY_NORMAL,
        reason="claim_download_lease",
    )

    await log_task_event(
        task,
        event_type="DOWNLOAD_START",
        from_stage=old_stage,
        to_stage=DOWNLOAD_STAGE_DOWNLOADING,
        detail={
            "attempt": task.download_attempt_count,
            "lease_until": task.download_lease_until,
            "download_celery_task_id": task.download_celery_task_id,
        },
    )
    return True


async def _mark_retry_waiting(db: AsyncSession, task: ChatGenerationTask, exc: Exception) -> datetime:
    now = _now()
    attempt = int(task.download_attempt_count or task.retry_count or 0)
    next_retry_at = _retry_at(attempt, now)
    error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc)

    retry_celery_task_id = _build_celery_task_id(task.id, attempt=attempt, reason="retry_waiting")

    task.pipeline_stage = DOWNLOAD_STAGE_RETRY_WAITING
    task.download_celery_task_id = retry_celery_task_id
    task.download_next_retry_at = next_retry_at
    task.download_lease_until = None
    task.download_last_error = error_message
    task.retry_count = attempt
    await db.commit()

    await _register_active_from_task(
        task,
        check_at=next_retry_at,
        priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
        reason="download_retry_waiting",
    )

    await log_task_event(
        task,
        event_type="DOWNLOAD_RETRY_WAITING",
        message=error_message,
        to_stage=DOWNLOAD_STAGE_RETRY_WAITING,
        detail={
            "attempt": attempt,
            "next_retry_at": next_retry_at,
            "download_celery_task_id": retry_celery_task_id,
        },
    )
    return next_retry_at


def _should_final_fail(task: ChatGenerationTask) -> bool:
    return int(task.download_attempt_count or task.retry_count or 0) >= int(settings.DOWNLOAD_TASK_MAX_ATTEMPTS or 3)


async def _run(task_id: str):
    async with async_session() as db:
        result = await db.execute(
            select(ChatGenerationTask).where(
                ChatGenerationTask.id == task_id,
                ChatGenerationTask.deleted_at.is_(None),
            ).with_for_update().limit(1)
        )
        task = result.scalar_one_or_none()
        if not task:
            return

        claimed = await _claim_download_lease(db, task)
        if not claimed:
            return

        try:
            downloaded = await download_generation_result(task)

            task = await _reload_task(db, task_id)
            if not task:
                return

            if task.gen_type == "image":
                task.image_url = downloaded.url
            else:
                task.video_url = downloaded.url
                task.video_cover_url = downloaded.cover_url

            task.status = "completed"
            task.pipeline_stage = DOWNLOAD_STAGE_DONE
            task.generated_at = _now()
            task.retry_count = 0
            task.download_lease_until = None
            task.download_next_retry_at = None
            task.download_last_error = None

            await record_chat_task_generated_resource(
                db,
                task,
                resource_url=downloaded.url,
                storage_path=downloaded.storage_path,
                file_size_bytes=downloaded.file_size_bytes,
                remote_url=task.remote_result_url,
                generated_at=task.generated_at,
            )

            await db.commit()
            await remove_download_active(task.id)

            await log_task_event(
                task,
                event_type="DOWNLOAD_SUCCESS",
                to_status="completed",
                to_stage=DOWNLOAD_STAGE_DONE,
                detail={
                    "resource_url": downloaded.url,
                    "video_cover_url": downloaded.cover_url,
                    "file_size_bytes": downloaded.file_size_bytes,
                    "download_attempt_count": task.download_attempt_count,
                },
            )

        except Exception as exc:
            try:
                await db.rollback()
            except Exception:
                pass

            task = await _reload_task(db, task_id)
            if not task:
                return

            if _should_final_fail(task):
                error_message = extract_error_message(exc, "下载") if callable(extract_error_message) else str(exc)
                await mark_chat_generation_task_failed_and_refund_once(
                    db,
                    task=task,
                    error_message=error_message,
                    pipeline_stage=DOWNLOAD_STAGE_FAILED,
                )
                task.download_last_error = error_message
                task.download_lease_until = None
                task.download_next_retry_at = None
                await db.commit()

                await remove_download_active(task.id)

                await log_task_event(
                    task,
                    event_type="DOWNLOAD_FAILED",
                    message=task.error_message,
                    detail={
                        "download_attempt_count": task.download_attempt_count,
                        "max_attempts": settings.DOWNLOAD_TASK_MAX_ATTEMPTS,
                    },
                )
            else:
                next_retry_at = await _mark_retry_waiting(db, task, exc)

                if celery_app:
                    delay_seconds = max(1, int((next_retry_at - _now()).total_seconds()))
                    download_generation_result_task.apply_async(
                        args=[task.id],
                        queue=DOWNLOAD_QUEUE,
                        priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
                        countdown=delay_seconds,
                        task_id=task.download_celery_task_id,
                    )


if celery_app:
    @celery_app.task(name="generation.download_generation_result_task", bind=True, max_retries=3, default_retry_delay=30)
    def download_generation_result_task(self, task_id: str):
        return run_async(_run(task_id))
else:
    class _DisabledTask:
        def delay(self, *args, **kwargs):
            raise RuntimeError("Celery is disabled")

        def apply_async(self, *args, **kwargs):
            raise RuntimeError("Celery is disabled")

    download_generation_result_task = _DisabledTask()