beige/backend/app/workers/tasks.py

"""
app/workers/tasks.py — Celery 任务入口（编排层）
只接收 task_id（基石B）。具体步骤在 pipeline_steps / pipeline_io。
打包任务在 packaging_task.py，此处 re-export 保持旧引用不变。
"""

import logging

from app.workers.celery_app import celery_app
# re-export 保持旧代码 `from app.workers.tasks import build_delivery_package` 不变
from app.workers.packaging_task import build_delivery_package  # noqa: F401
from app.workers.replenish_task import replenish_text_candidates  # noqa: F401

logger = logging.getLogger(__name__)


def _get_db():
    """获取同步 DB Session（Celery worker 环境用同步）"""
    from app.core.database import SessionLocal
    return SessionLocal()


def _push_event_sync(task_id: int, workspace_id: int, event: str, data: dict, seq: int):
    """同步推 SSE 事件（worker 内部用）"""
    import json as _json
    import redis as redis_lib
    from app.core.config import get_settings
    s = get_settings()
    r = redis_lib.from_url(s.REDIS_URL, decode_responses=True)
    hist_key = f"sse:task:{task_id}:events"
    channel = f"sse:task:{task_id}"
    record = _json.dumps({"event": event, "data": data, "seq": seq,
                          "workspace_id": workspace_id}, ensure_ascii=False)
    r.rpush(hist_key, record)
    r.expire(hist_key, 3600)
    r.publish(channel, record)


@celery_app.task(
    bind=True,
    name="app.workers.tasks.run_generation_pipeline",
    max_retries=3,
    default_retry_delay=30,
    queue="generation",
)
def run_generation_pipeline(self, task_id: int) -> dict:
    """
    生产链主任务。只接收 task_id，绝不接收 key。
    子步骤委托给 pipeline_steps（查库/解密/飞轮上下文）
    和 pipeline_io（文案/图片生成+存储）。
    """
    logger.info("run_generation_pipeline start: task_id=%s", task_id)
    db = _get_db()
    seq = 0
    workspace_id = 0  # G3坑修复：确保 except 里能用真实 workspace_id

    try:
        from app.constants.enums import TaskStatus
        from app.workers.pipeline_steps import (
            load_task_and_product,
            decrypt_user_key,
            build_clients_and_clear_key,
            build_product_dict,
            load_flywheel_context,
        )
        from app.workers.pipeline_io import run_text_generation, run_image_generation
        from app.core.config import get_settings

        # Step1: 查任务 + 产品
        task, product = load_task_and_product(db, task_id)
        if not task:
            return {"task_id": task_id, "status": "not_found"}

        workspace_id = task.workspace_id

        # Step2: Fernet 解密（局部变量，不传出）
        plain_key = decrypt_user_key(db, task.operator_id, workspace_id)

        # Step3: 构建 AIClients
        clients = build_clients_and_clear_key(plain_key)
        plain_key = None  # 用完即清除，基石B

        task.status = TaskStatus.GENERATING
        db.commit()

        # Step4: 推 task_started + 飞轮上下文
        seq += 1
        _push_event_sync(task_id, workspace_id, "task_started", {
            "task_id": task_id, "total_text": task.text_count, "total_image": task.image_count,
        }, seq)

        product_dict = build_product_dict(product)
        flywheel_fragment, flywheel_ctx = load_flywheel_context(db, workspace_id, task.product_id, product_dict)

        if flywheel_fragment:
            seq += 1
            _push_event_sync(task_id, workspace_id, "flywheel_injected", {
                "recent_preference": flywheel_ctx.get("recent_preference", ""),
                "reject_reasons": flywheel_ctx.get("reject_reasons", []),
            }, seq)

        # Step5: 文案生成（S1: 返回 needs_replenish=合格数<目标数，触发后台补充）
        candidates_raw, seq, needs_replenish = run_text_generation(
            db, clients, task, product_dict, flywheel_fragment,
            _push_event_sync, workspace_id, seq,
        )
        if needs_replenish:
            # 合格文案不足用户目标条数：先展示已合格的，后台异步补充（先展示后台补铁律）
            # 注：candidates_raw 是本轮原始生成数(含被过滤的低分条)，非合格入库数；
            # 真实合格数以库内 saved_count 为准，补充子任务会按库内实际缺口补到够。
            logger.warning(
                "文案合格数不足将后台补充: task_id=%s 本轮生成=%d 目标=%d",
                task_id, len(candidates_raw), task.text_count,
            )
            try:
                replenish_text_candidates.delay(task_id)
            except Exception as _re:
                logger.error("触发后台补充失败: task_id=%s err=%s", task_id, _re)

        # Step6+7+8: 图片生成 + 后处理 + 存盘
        first_copy = candidates_raw[0] if candidates_raw else {}
        settings = get_settings()
        seq = run_image_generation(
            db, clients, task, product_dict,
            _push_event_sync, workspace_id, seq,
            first_copy, settings.UPLOAD_BASE_PATH,
        )

        # 最终状态 + task_done
        task.status = TaskStatus.PENDING_SELECTION
        db.commit()

        seq += 1
        _push_event_sync(task_id, workspace_id, "task_done", {
            "task_id": task_id, "status": "pending_selection"
        }, seq)

        logger.info("run_generation_pipeline done: task_id=%s", task_id)
        return {"task_id": task_id, "status": "pending_selection"}

    except Exception as exc:
        logger.error("run_generation_pipeline failed: task_id=%s err=%s", task_id, exc)
        # 重试耗尽才落 FAILED 终态;还能重试则保持 GENERATING(别伪装成"待开始")。
        # B6:失败要可见,区分"没跑(pending)"和"跑挂了(failed)"。
        exhausted = self.request.retries >= self.max_retries
        try:
            from app.models.task import GenerationTask as GT
            from app.constants.enums import TaskStatus as TS
            t = db.query(GT).filter(GT.id == task_id).first()
            if t:
                t.status = TS.FAILED if exhausted else TS.GENERATING
                db.commit()
        except Exception:
            pass
        _push_event_sync(task_id, workspace_id, "error", {"code": 50001, "message": str(exc)}, seq + 1)
        if exhausted:
            return {"task_id": task_id, "status": "failed", "error": str(exc)}
        raise self.retry(exc=exc)
    finally:
        db.close()