xiaoxia-saas/apps/api/app/tasks/lipsync_tts.py

"""AI 数字人对口型 TTS 异步任务 — 将 TTS 合成从 HTTP 请求移至 Celery 后台执行.

优化目标：将 create_job 的 API 响应时间从 6~35s 降到 <1s。
任务流程：
1. 创建新 DB session，加载 job 记录
2. 调用 CosyVoice 合成音频
3. 下载音频并转存到自家 OSS
4. 更新 job 的 audio_url
5. 签名 URL 并提交到 MediaKit
6. 更新 job 状态为 submitted
7. 异常时标记 job 为 failed

注意：使用 @shared_task 而非绑定到某个 celery_app 实例，
确保任务能被 Worker 侧 celery_app 正确注册，同时 API 侧 send_task/apply_async 仍可正常调用。
"""

import io
import logging
from datetime import datetime, timezone
from urllib.parse import urlparse

from celery import shared_task

logger = logging.getLogger(__name__)

# MediaKit 预签名 URL 有效期（7天，秒），与 LipsyncService._sign_media_url 保持一致
_MEDIAKIT_URL_TTL_SECONDS = 7 * 24 * 3600


def _sign_media_url(url: str) -> str:
    """对自家 OSS 私有桶 URL 重签长有效期预签名.

    - 自家 OSS URL → 重签 7 天有效期
    - 外部临时 URL → 原样透传
    - 任何异常降级原样返回，不阻断主流程
    """
    if not url:
        return url
    try:
        from packages.shared.storage import get_shared_storage_service

        storage = get_shared_storage_service()
        public_base = getattr(storage, "public_url", "")
        if not isinstance(public_base, str) or not public_base:
            return url
        own_host = urlparse(public_base).netloc.lower()
        host = urlparse(url).netloc.lower()
        if not own_host or host != own_host:
            return url
        signed = storage.get_download_url(url, expires_seconds=_MEDIAKIT_URL_TTL_SECONDS)
        return signed or url
    except Exception as exc:  # noqa: BLE001
        logger.warning("[lipsync_tts] URL 重签失败，原样返回: url_prefix=%s err=%s", url[:80], exc)
        return url


@shared_task(
    bind=True,
    name="lipsync_tts.synthesize_and_submit",
    max_retries=2,
    default_retry_delay=30,
)
def tts_synthesize_and_submit(
    self,
    job_id: str,
    user_id: str,
    voice_id: str,
    script_text: str,
    speed: float,
    emotion: str,
):
    """异步执行 TTS 合成 + OSS 转存 + MediaKit 提交.

    在 Celery worker 中运行，不阻塞 HTTP 请求。
    """
    from app.services.mediakit_client import MediaKitError, get_mediakit_client
    from sqlalchemy.orm import Session as DBSession

    from packages.adapters.sqlalchemy_impl.models import LipsyncJobModel
    from packages.application.cosyvoice_service import CosyVoiceError, CosyVoiceService
    from packages.shared.url_security import safe_download_bytes

    # SessionLocal 获取：
    # - API 容器：app.db.SessionLocal（环境变量完整，导入即建引擎）
    # - Worker 容器：worker_app.db.SessionLocal（Worker 自己的 settings 初始化引擎）
    # API 侧没有 worker_app 模块 → ImportError 直接回退；
    # Worker 侧 app.db 会因缺少 API 专有环境变量抛 pydantic ValidationError，
    # 此时也要回退到 worker_app.db。
    try:
        from worker_app.db import SessionLocal  # type: ignore
    except Exception:  # noqa: BLE001
        from app.db import SessionLocal  # type: ignore

    db: DBSession = SessionLocal()
    try:
        job = (
            db.query(LipsyncJobModel)
            .filter(
                LipsyncJobModel.id == job_id,
                LipsyncJobModel.user_id == user_id,
            )
            .first()
        )

        if job is None:
            logger.error("[lipsync_tts] Job not found: job_id=%s", job_id)
            return

        # 已取消的任务不再处理
        if job.status == "cancelled":
            logger.info("[lipsync_tts] Job already cancelled, skipping: job_id=%s", job_id)
            return

        # 1. TTS 合成
        try:
            cosyvoice = CosyVoiceService()
            result = cosyvoice.submit_synthesize_task(
                text=script_text,
                voice_id=voice_id,
                speed=speed,
                emotion=emotion,
            )
        except CosyVoiceError as exc:
            logger.error("[lipsync_tts] TTS 合成失败: job_id=%s err=%s", job_id, exc)
            job.status = "failed"
            job.error_message = f"TTS 合成失败: {exc}"
            job.error_code = "TTSSynthesisFailed"
            job.updated_at = datetime.now(timezone.utc)
            db.commit()
            return
        except ValueError as exc:
            logger.error("[lipsync_tts] TTS 参数错误: job_id=%s err=%s", job_id, exc)
            job.status = "failed"
            job.error_message = f"TTS 参数错误: {exc}"
            job.error_code = "TTSInvalidParam"
            job.updated_at = datetime.now(timezone.utc)
            db.commit()
            return

        temp_url = result.get("audio_url", "")
        if not temp_url:
            logger.error("[lipsync_tts] TTS 未返回音频 URL: job_id=%s", job_id)
            job.status = "failed"
            job.error_message = "TTS 未返回音频 URL"
            job.error_code = "TTSNoAudio"
            job.updated_at = datetime.now(timezone.utc)
            db.commit()
            return

        # 2. 下载并转存到自家 OSS
        try:
            audio_data = safe_download_bytes(
                temp_url,
                purpose="lipsync_tts_audio",
                allowed_mime_types=(
                    "audio/mpeg",
                    "audio/mp3",
                    "audio/wav",
                    "audio/x-wav",  # CosyVoice 部分接口返回 audio/x-wav，与 audio/wav 等价（RIFF/WAVE）
                    "audio/mp4",
                    "audio/x-m4a",
                ),
                timeout=60.0,
            )
            from packages.shared.storage import get_shared_storage_service

            storage = get_shared_storage_service()
            storage_key = f"lipsync-tts/{user_id}/{job_id}.mp3"
            permanent_url = storage.upload_file(io.BytesIO(audio_data), storage_key, content_type="audio/mpeg")
            logger.info("[lipsync_tts] TTS 音频已转存 OSS: job_id=%s key=%s", job_id, storage_key)
            job.audio_url = permanent_url
        except Exception as exc:
            logger.warning(
                "[lipsync_tts] TTS 音频转存 OSS 失败，回退临时 URL: job_id=%s err=%s",
                job_id,
                exc,
            )
            job.audio_url = temp_url

        db.commit()

        # 3. 签名 URL 并提交到 MediaKit（复用模块内 _sign_media_url，避免对 LipsyncService 的耦合）
        audio_url = _sign_media_url(job.audio_url)
        video_url = _sign_media_url(job.video_url)

        client = get_mediakit_client()
        try:
            mk_result = client.submit_lipsync(
                video_url=video_url,
                audio_url=audio_url,
                enable_video_loop=job.enable_video_loop,
                client_token=job_id,
            )
            job.mediakit_task_id = mk_result["task_id"]
            job.status = "submitted"
            job.submitted_at = datetime.now(timezone.utc)
            logger.info(
                "[lipsync_tts] 已提交 MediaKit: job_id=%s task_id=%s",
                job_id,
                mk_result["task_id"],
            )
        except MediaKitError as exc:
            job.status = "failed"
            job.error_message = str(exc)
            job.error_code = exc.code
            logger.error("[lipsync_tts] 提交 MediaKit 失败: job_id=%s err=%s", job_id, exc)

        db.commit()

        # 4. 链式触发 Celery 兜底轮询：MediaKit 提交成功后由 worker 主动拉取状态到终态，
        # 不依赖前端轮询触发的 FastAPI background task（后台任务可能静默失败导致永久卡 running）
        if job.status == "submitted" and job.mediakit_task_id:
            poll_mediakit_status.apply_async(
                kwargs={"job_id": job_id, "user_id": user_id},
                countdown=10,  # 10 秒后开始轮询（给 MediaKit 一点处理时间）
            )

    except Exception:
        logger.exception("[lipsync_tts] 未预期的异常: job_id=%s", job_id)
        try:
            job = db.query(LipsyncJobModel).filter(LipsyncJobModel.id == job_id).first()
            if job and job.status not in ("cancelled", "failed", "completed"):
                job.status = "failed"
                job.error_message = "TTS 异步任务执行异常"
                job.error_code = "AsyncTaskError"
                job.updated_at = datetime.now(timezone.utc)
                db.commit()
        except Exception:
            logger.exception("[lipsync_tts] 回写失败状态时异常: job_id=%s", job_id)
    finally:
        db.close()


@shared_task(
    bind=True,
    name="lipsync_tts.poll_mediakit_status",
    max_retries=60,  # 最多轮询 60 次
    default_retry_delay=10,  # 每次间隔 10 秒（总兜底时长 10 分钟）
)
def poll_mediakit_status(self, job_id: str, user_id: str):
    """Celery 兜底轮询：TTS 提交 MediaKit 后，由 worker 主动拉取状态直到终态。

    不依赖前端轮询，避免 background task 静默失败导致任务永久卡 running/submitted。
    """
    from sqlalchemy.orm import Session as DBSession

    from packages.adapters.sqlalchemy_impl.models import LipsyncJobModel

    try:
        from worker_app.db import SessionLocal  # type: ignore
    except Exception:  # noqa: BLE001
        from app.db import SessionLocal  # type: ignore

    db: DBSession = SessionLocal()
    try:
        job = (
            db.query(LipsyncJobModel)
            .filter(LipsyncJobModel.id == job_id, LipsyncJobModel.user_id == user_id)
            .first()
        )
        if job is None:
            logger.warning("[lipsync_poll] Job not found: job_id=%s", job_id)
            return

        # 已终态，不需要再轮询
        if job.status in ("completed", "failed", "cancelled"):
            return

        if not job.mediakit_task_id:
            logger.warning("[lipsync_poll] Job has no mediakit_task_id: job_id=%s status=%s", job_id, job.status)
            return

        from app.services.mediakit_client import MediaKitError, get_mediakit_client

        client = get_mediakit_client()
        try:
            status_data = client.get_task_status(job.mediakit_task_id)
        except MediaKitError as exc:
            logger.warning("[lipsync_poll] 拉取 MediaKit 状态失败，将重试: job_id=%s err=%s", job_id, exc)
            raise self.retry(exc=exc)

        mk_status = status_data.get("status", "running")

        if mk_status == "succeeded":
            # 复用 LipsyncService 的持久化逻辑
            from app.services.lipsync_service import LipsyncService

            svc = LipsyncService(db)
            result = status_data.get("result", {})
            job.status = "completed"
            output_url = result.get("video_url", "")
            job.output_video_url = svc._persist_output_video(output_url, job_id, user_id)
            job.output_duration = result.get("duration", 0.0)
            job.completed_at = datetime.now(timezone.utc)
            job.updated_at = datetime.now(timezone.utc)
            db.commit()
            logger.info("[lipsync_poll] 任务完成: job_id=%s", job_id)
        elif mk_status in ("failed", "error"):
            error = status_data.get("error", {})
            job.status = "failed"
            job.error_message = error.get("message", "任务执行失败")
            job.error_code = error.get("code", "TaskFailed")
            job.completed_at = datetime.now(timezone.utc)
            job.updated_at = datetime.now(timezone.utc)
            db.commit()
            logger.info("[lipsync_poll] 任务失败: job_id=%s err=%s", job_id, job.error_message)
        else:
            # 中间状态，更新时间戳，继续重试
            job.updated_at = datetime.now(timezone.utc)
            if isinstance(mk_status, str) and mk_status:
                job.status = mk_status
            db.commit()
            logger.debug("[lipsync_poll] 任务仍在 %s，继续轮询: job_id=%s", mk_status, job_id)
            raise self.retry()
    except Exception as exc:
        logger.exception("[lipsync_poll] 未预期异常: job_id=%s", job_id)
        db.rollback()
        raise self.retry(exc=exc)
    finally:
        db.close()