SupplyAgent/supply_infra/pipeline/aigc_outbox.py

from __future__ import annotations

import hashlib
import json
import logging
from typing import Any

from supply_infra.aigc.client import AigcClient
from supply_infra.db.repositories.pipeline_outbox_repo import PipelineOutboxRepository
from supply_infra.db.session import get_session
from supply_infra.pipeline.contracts import StepContext
from supply_infra.pipeline.dates import china_now
from supply_infra.pipeline.enums import OutboxStatus
from supply_infra.scheduler.jobs.publish_videos_from_discovery import (
    prepare_publish_batches,
)
from supply_infra.services.video_discovery_service import (
    get_video_discovery_service,
)

logger = logging.getLogger(__name__)

_EFFECT_TYPE = "aigc_publish_batch"
_CLAIMABLE_STATUSES = {
    OutboxStatus.PENDING.value,
    OutboxStatus.CREATED.value,
    OutboxStatus.RETRYABLE_FAILED.value,
}


def _canonical_payload(payload: dict[str, Any]) -> tuple[str, str]:
    canonical = json.dumps(
        payload,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        default=str,
    )
    return canonical, hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _enqueue_batches(
    context: StepContext,
    batches: list[dict[str, Any]],
) -> list[str]:
    outbox_ids: list[str] = []
    with get_session() as session:
        repo = PipelineOutboxRepository(session)
        for payload in batches:
            _, payload_hash = _canonical_payload(payload)
            batch_key = str(payload["batch_key"])
            mode = "dry_run" if context.dry_run else "live"
            idempotency_key = (
                f"aigc_publish:{context.biz_dt}:{mode}:{batch_key}"
            )
            row = repo.enqueue(
                run_id=context.run_id,
                step_run_id=context.step_run_id,
                effect_type=_EFFECT_TYPE,
                idempotency_key=idempotency_key,
                payload_hash=payload_hash,
                payload=payload,
                dry_run=context.dry_run,
            )
            outbox_ids.append(str(row.outbox_id))
    return outbox_ids


def _mark_ambiguous(outbox_id: str, error: str) -> dict[str, Any]:
    with get_session() as session:
        repo = PipelineOutboxRepository(session)
        row = repo.get(outbox_id, for_update=True)
        if row is None:
            raise ValueError(f"outbox event not found: {outbox_id}")
        repo.mark_ambiguous(row, error=error)
    return {
        "success": False,
        "outbox_id": outbox_id,
        "status": OutboxStatus.AMBIGUOUS.value,
        "error_code": "aigc_outbox_ambiguous",
        "error": error,
    }


def _mark_retryable(
    outbox_id: str,
    error: str,
    response: dict[str, Any] | None = None,
) -> dict[str, Any]:
    with get_session() as session:
        repo = PipelineOutboxRepository(session)
        row = repo.get(outbox_id, for_update=True)
        if row is None:
            raise ValueError(f"outbox event not found: {outbox_id}")
        repo.mark_retryable_failed(row, error=error, response=response)
    return {
        "success": False,
        "outbox_id": outbox_id,
        "status": OutboxStatus.RETRYABLE_FAILED.value,
        "error_code": "aigc_outbox_retryable",
        "error": error,
    }


def dispatch_aigc_outbox_event(
    outbox_id: str,
    *,
    client: AigcClient | None = None,
) -> dict[str, Any]:
    """
    Dispatch one durable event.

    A stale ``sending`` event is deliberately moved to ``ambiguous`` instead
    of being replayed: the remote create request may have succeeded before the
    process died. This trades automatic recovery for duplicate-action safety.
    """
    with get_session() as session:
        repo = PipelineOutboxRepository(session)
        row = repo.get(outbox_id, for_update=True)
        if row is None:
            raise ValueError(f"outbox event not found: {outbox_id}")
        if row.dry_run:
            repo.mark_succeeded(
                row,
                response={
                    "dry_run": True,
                    "external_request_made": False,
                    "message": "AIGC dispatch intentionally skipped",
                },
            )
            return {
                "success": True,
                "outbox_id": outbox_id,
                "status": OutboxStatus.SUCCEEDED.value,
                "dry_run": True,
                "external_request_made": False,
            }
        if row.status == OutboxStatus.SUCCEEDED.value:
            return {
                "success": True,
                "outbox_id": outbox_id,
                "status": row.status,
                "skipped_already_succeeded": True,
                "external_id": row.external_id,
            }
        if row.status == OutboxStatus.AMBIGUOUS.value:
            return {
                "success": False,
                "outbox_id": outbox_id,
                "status": row.status,
                "error_code": "aigc_outbox_ambiguous",
                "error": row.record_error or "外部请求结果不确定，需要人工核对",
            }
        if row.status == OutboxStatus.SENDING.value:
            repo.mark_ambiguous(
                row,
                error=(
                    "检测到未确认的发送中事件；为避免重复创建 AIGC 计划，"
                    "已停止自动重放，请按确定性计划名人工核对"
                ),
            )
            return {
                "success": False,
                "outbox_id": outbox_id,
                "status": OutboxStatus.AMBIGUOUS.value,
                "error_code": "aigc_outbox_ambiguous",
                "error": row.record_error,
            }
        if row.status not in _CLAIMABLE_STATUSES:
            return {
                "success": False,
                "outbox_id": outbox_id,
                "status": row.status,
                "error_code": "aigc_outbox_not_dispatchable",
                "error": f"事件状态不允许发送: {row.status}",
            }

        payload = dict(row.payload_json or {})
        external_id = str(row.external_id or "")
        external_name = str(row.external_name or "")
        repo.mark_sending(row, now=china_now())

    required = {
        "aweme_ids",
        "candidate_ids",
        "plan_name",
        "plan_label",
        "produce_plan_id",
        "publish_plan_id",
    }
    missing = sorted(required - payload.keys())
    if missing:
        return _mark_retryable(
            outbox_id,
            f"Outbox payload 缺少字段: {', '.join(missing)}",
        )

    active_client = client or AigcClient()
    if not external_id:
        try:
            create_result = active_client.create_video_crawler_plan(
                [str(value) for value in payload["aweme_ids"]],
                plan_name=str(payload["plan_name"]),
            )
        except Exception as exc:
            logger.exception("AIGC create request outcome is ambiguous: outbox=%s", outbox_id)
            return _mark_ambiguous(
                outbox_id,
                f"创建请求异常，远端是否成功未知: {exc}",
            )

        if not create_result.get("success"):
            return _mark_retryable(
                outbox_id,
                str(create_result.get("error") or "创建爬取计划失败"),
                response=create_result,
            )
        external_id = str(create_result.get("crawler_plan_id") or "")
        external_name = str(
            create_result.get("crawler_plan_name") or payload["plan_name"]
        )
        if not external_id:
            return _mark_ambiguous(
                outbox_id,
                "AIGC 返回成功但未提供 crawler_plan_id，无法安全重试",
            )
        with get_session() as session:
            repo = PipelineOutboxRepository(session)
            row = repo.get(outbox_id, for_update=True)
            if row is None:
                raise ValueError(f"outbox event not found: {outbox_id}")
            repo.mark_created(
                row,
                external_id=external_id,
                external_name=external_name,
                response=create_result,
            )

    try:
        bind_result = active_client.bind_crawler_to_produce_plan(
            external_id,
            str(payload["produce_plan_id"]),
            crawler_plan_name=external_name or str(payload["plan_name"]),
        )
    except Exception as exc:
        logger.exception("AIGC bind failed after create: outbox=%s", outbox_id)
        return _mark_retryable(outbox_id, f"绑定生成计划异常: {exc}")

    if not bind_result.get("success"):
        return _mark_retryable(
            outbox_id,
            str(bind_result.get("error") or "绑定生成计划失败"),
            response=bind_result,
        )

    try:
        updated = get_video_discovery_service().mark_candidates_aigc_plans(
            [int(value) for value in payload["candidate_ids"]],
            crawler_plan_id=external_id,
            produce_plan_id=str(payload["produce_plan_id"]),
            publish_plan_id=str(payload["publish_plan_id"]),
            plan_label=str(payload["plan_label"]),
        )
    except Exception as exc:
        logger.exception("Failed to persist AIGC result: outbox=%s", outbox_id)
        return _mark_retryable(outbox_id, f"本地保存 AIGC 结果失败: {exc}")

    final_response = {
        "create": {
            "crawler_plan_id": external_id,
            "crawler_plan_name": external_name,
        },
        "bind": bind_result,
        "candidate_rows_updated": updated,
    }
    with get_session() as session:
        repo = PipelineOutboxRepository(session)
        row = repo.get(outbox_id, for_update=True)
        if row is None:
            raise ValueError(f"outbox event not found: {outbox_id}")
        repo.mark_succeeded(row, response=final_response)

    return {
        "success": True,
        "outbox_id": outbox_id,
        "status": OutboxStatus.SUCCEEDED.value,
        "external_id": external_id,
        "candidate_rows_updated": updated,
    }


def execute_aigc_outbox(context: StepContext) -> dict[str, Any]:
    prepared = prepare_publish_batches(biz_dt=context.biz_dt)
    outbox_ids = _enqueue_batches(context, list(prepared["batches"]))

    # A retry can happen after candidate rows were marked but before the event
    # was acknowledged. Include all nonterminal events belonging to this run.
    with get_session() as session:
        for row in PipelineOutboxRepository(session).list_for_run(context.run_id):
            if (
                row.effect_type == _EFFECT_TYPE
                and row.status != OutboxStatus.SUCCEEDED.value
            ):
                outbox_ids.append(str(row.outbox_id))

    ordered_ids = list(dict.fromkeys(outbox_ids))
    results = [dispatch_aigc_outbox_event(outbox_id) for outbox_id in ordered_ids]
    failures = [result for result in results if result.get("success") is False]
    return {
        "success": not failures,
        "effect_recorded": bool(ordered_ids),
        "external_request_made": bool(ordered_ids) and not context.dry_run,
        "dry_run": context.dry_run,
        "biz_dt": context.biz_dt,
        "candidate_count": prepared["candidate_count"],
        "batch_count": len(ordered_ids),
        "failed_batch_count": len(failures),
        "outbox_ids": ordered_ids,
        "batches": results,
        "error_code": (
            str(failures[0].get("error_code") or "aigc_publish_failed")
            if failures
            else None
        ),
        "error": (
            str(failures[0].get("error") or "AIGC Outbox 发送失败")
            if failures
            else None
        ),
    }