video-gen/video-gen-api/app/services/generation/pipeline/lifecycle_service.py

from __future__ import annotations

import asyncio
from datetime import datetime, timedelta, timezone

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

from app.config import settings
from app.enums.generation_status import GenerationRecordPipelineStage
from app.enums.generation_task import ChatGenerationPipelineStage, GenerationType
from app.models.chat_generation_task import ChatGenerationTask
from app.models.generation_record import GenerationRecord
from app.models.base import async_session
from app.services.generation.pipeline.db_lock_service import DatabaseRowLockBusy
from app.services.generation.pipeline.owner_service import GenerationOwner
from app.services.generation.refund_service import (
    mark_chat_generation_task_failed_and_refund_once,
    mark_generation_record_failed_and_refund_once,
)


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


def generation_deadline(*, gen_type: str, started_at: datetime) -> datetime:
    if str(gen_type or "").lower() == GenerationType.IMAGE.value:
        return started_at + timedelta(minutes=max(1, int(settings.CHATAPI_ASYNC_IMAGE_DEADLINE_MINUTES or 30)))
    return started_at + timedelta(hours=max(1, int(settings.CHATAPI_ASYNC_VIDEO_FINAL_DEADLINE_HOURS or 24)))


def reset_execution_fields(owner: GenerationOwner, *, started_at: datetime, attempt_no: int) -> None:
    owner.generation_attempt_no = max(1, int(attempt_no or 1))
    owner.resource_generation_started_at = started_at
    owner.deadline_at = generation_deadline(gen_type=owner.gen_type, started_at=started_at)
    owner.error_message = None
    owner.seedance_task_id = None
    if hasattr(owner, "provider_task_id"):
        owner.provider_task_id = None
    owner.remote_result_url = None
    owner.provider_response_json = None
    owner.provider_create_claim_token = None
    owner.provider_create_lease_until = None
    owner.provider_create_started_at = None
    owner.retry_count = int(getattr(owner, "manual_retry_count", 0) or 0)
    owner.poll_error_count = 0
    owner.poll_count = 0
    owner.last_poll_at = None
    owner.poll_started_at = None
    owner.next_poll_at = None
    owner.poll_interval_seconds = 0
    owner.poll_claim_token = None
    owner.poll_lease_until = None
    owner.download_celery_task_id = None
    owner.download_enqueued_at = None
    owner.download_started_at = None
    owner.download_claim_token = None
    owner.download_lease_until = None
    owner.download_next_retry_at = None
    owner.download_attempt_count = 0
    owner.download_last_error = None
    owner.download_storage_date_dir = None
    owner.generated_at = None
    owner.image_url = None
    owner.video_url = None
    owner.video_cover_url = None


async def notify_owner_finished(db: AsyncSession, owner: GenerationOwner) -> None:
    """Run module hooks after the generation owner transaction has committed.

    Module project/step rows have their own short locks. A transient lock conflict
    must not turn an already completed generation into a download/provider retry,
    so hooks use a fresh short transaction with a small local retry window.
    """
    _ = db
    if not isinstance(owner, ChatGenerationTask):
        return
    from app.services.generation.module_hook_service import notify_chat_generation_task_finished
    from app.services.generation.ai.task_group_service import aggregate_parent_for_child
    from app.services.operation_log_service import log_operation_event

    task_id = str(owner.id)
    attempt_no = int(owner.generation_attempt_no or 1)
    last_error: Exception | None = None
    for retry_index in range(3):
        try:
            async with async_session() as hook_db:
                result = await hook_db.execute(
                    select(ChatGenerationTask).where(
                        ChatGenerationTask.id == task_id,
                        ChatGenerationTask.deleted_at.is_(None),
                    ).limit(1)
                )
                fresh_task = result.scalar_one_or_none()
                if fresh_task is None:
                    return
                await notify_chat_generation_task_finished(hook_db, fresh_task)
                await aggregate_parent_for_child(hook_db, fresh_task)
                await hook_db.commit()
                return
        except DatabaseRowLockBusy as exc:
            last_error = exc
            await asyncio.sleep(1 + retry_index)
        except Exception as exc:
            last_error = exc
            break

    log_operation_event(
        domain="generation_pipeline",
        event_type="MODULE_HOOK_DEFERRED_MANUAL_CHECK",
        event_status="failed",
        source="pipeline",
        task_id=task_id,
        message="生成任务已进入终态，但模块状态回填失败，需要人工排查",
        detail={
            "generation_attempt_no": attempt_no,
            "retry_count": 3,
        },
        error=str(last_error or "unknown module hook error"),
    )


async def mark_owner_failed_and_refund_once(
    db: AsyncSession,
    owner: GenerationOwner,
    *,
    error_message: str,
    pipeline_stage: str,
) -> GenerationOwner:
    if isinstance(owner, ChatGenerationTask):
        await mark_chat_generation_task_failed_and_refund_once(
            db,
            task=owner,
            error_message=error_message,
            pipeline_stage=pipeline_stage,
            generation_attempt_no=int(owner.generation_attempt_no or 1),
        )
        return owner

    owner.pipeline_stage = pipeline_stage if pipeline_stage in {item.value for item in GenerationRecordPipelineStage} else GenerationRecordPipelineStage.FAILED.value
    await mark_generation_record_failed_and_refund_once(
        db,
        record=owner,
        error_message=error_message,
        generation_attempt_no=int(owner.generation_attempt_no or 1),
    )
    return owner