patbond-ai/app/providers/openai/tasks.py

"""同步型供应商的任务注册表：把「一次调用等到底」转成「提交 + 查询」。

## 为什么必须有它

ADR-023 §4 定的是「Java↔Python 边界一律异步，同步型供应商由其适配层在内部消化
这个差异」。火山的视频接口本身异步，所以视频适配器不需要这一层。**OpenAI 的图片
接口是同步的**——`POST /images/generations` 在一次请求里等到图片生成完才返回，
没有任务 ID，也没有查询接口。差异必须在这里被吸收，否则就会渗透到客户端契约、
等待态设计和 E2E，那是 §4 明确要避免的。

## 这与「Python 不持有状态」的张力，以及取舍

Seedance 那条路径上，Python 确实无状态：任务状态的唯一事实来源是 Java 的
`generation_jobs`，Python 只是个转发器。**图片这条路径做不到**——同步接口没有可
供事后查询的厂商 ID，所以「在途任务」这件事只能记在本进程内存里。

代价与处置如实登记：

1. **进程重启会丢失在途任务。** 之后 Java 轮询会得到 `task_not_found`，任务被判
   终态失败，而**那次调用可能已经完成并计费**。OpenAI 的图片接口没有任务 ID 可供
   事后查询，所以这个窗口**无法靠"再查一次"弥补**，只能由 `max_attempts` 兜住。
   这与火山那边「提交超时不知道有没有扣费」是同一类风险，但触发面更大（那边只有
   提交超时，这边还包括进程重启）。
2. **产物是 base64 字节，只在内存里。** 完成的任务必须在 TTL 内被取走，否则丢失。
   而转存阻塞于对象存储决策（ADR-023 §7），所以**图片链路在存储决策落地前比视频
   更受限**：视频至少有个 24 小时的厂商 URL 可以先看效果，图片的字节一旦被清理
   就什么都没有了。
3. 注册表**有界且有 TTL**，否则它就是个内存泄漏。

## 为什么用线程而不是 asyncio

FastAPI 的端点在本服务里是同步函数（`def` 而非 `async def`），且 `httpx.Client`
是同步客户端。用线程池承载那个长调用与现有代码一致；改成 asyncio 需要把传输层、
端点、测试全部改成异步，收益只是少几个线程——本服务的并发量级（一台轻量服务器上
的少量生成请求）用不上那个收益。
"""

from __future__ import annotations

import logging
import threading
import time
import uuid
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field

from app.domain.generation import Failure, GenerationResult, TaskState

log = logging.getLogger(__name__)

#: 本地任务 ID 的前缀，便于在日志里与厂商 ID 区分开。
LOCAL_TASK_PREFIX = "local-"


@dataclass
class _Entry:
    """一条在途或已完成的任务。"""

    future: Future
    created_at: float = field(default_factory=time.monotonic)
    #: **首次观察到**完成的时刻；未完成为 None。TTL 从这里起算而不是从创建起算——
    #: 否则一个耗时接近 TTL 的任务会在刚完成时就被清理掉，结果还没被取走。
    #:
    #: 为什么是「首次观察到」而不是「实际完成」：实际完成时刻需要一个
    #: `add_done_callback` 才能捕捉，而那会**死锁**——快速完成的任务在
    #: `add_done_callback` 时已经 done，回调便**同步地在调用线程执行**，而该线程
    #: 此刻正持有本注册表的锁（非可重入）。改用惰性记录既无此风险，又只会让结果
    #: 多留一会儿（保守方向），不会让它提前消失。
    finished_at: float | None = None


class SyncTaskRegistry:
    """把同步调用包装成可轮询的任务。

    线程安全：`submit`/`fetch`/清理可能来自不同请求线程。
    """

    def __init__(
        self,
        *,
        max_workers: int = 4,
        completed_ttl_sec: float = 900.0,
        max_entries: int = 256,
    ):
        """
        :param max_workers: 并发执行的同步调用数上限。**同时也是对服务商并发的
            自我限制**：超出的提交会排队而不是一起打过去，既避免把自己的 RPM 打爆，
            也让「同时有多少钱在花」有个上界。
        :param completed_ttl_sec: 已完成任务的保留时长。需大于 Java 侧最长的轮询
            间隔（当前阶梯最长 30 秒），否则结果会在被取走之前被清理。默认 15
            分钟留足余量。
        :param max_entries: 注册表容量上限。满了先清理过期项，仍满则拒绝新提交
            （而不是无声挤掉某个在途任务——那会让一个已经在计费的调用变成孤儿）。
        """
        self._executor = ThreadPoolExecutor(
            max_workers=max_workers, thread_name_prefix="sync-gen"
        )
        self._completed_ttl_sec = completed_ttl_sec
        self._max_entries = max_entries
        self._entries: dict[str, _Entry] = {}
        self._lock = threading.Lock()

    def submit(self, work: Callable[[], GenerationResult]) -> str:
        """安排一次同步调用在后台执行，立刻返回本地任务 ID。

        :raises RuntimeError: 注册表已满（并发在途任务过多）
        """
        with self._lock:
            self._evict_expired_locked()
            if len(self._entries) >= self._max_entries:
                raise RuntimeError(
                    f"在途生成任务已达上限 {self._max_entries}，请稍后重试"
                )
            task_id = LOCAL_TASK_PREFIX + uuid.uuid4().hex
            # 刻意不注册 done_callback，原因见 _Entry.finished_at 的说明（会死锁）。
            self._entries[task_id] = _Entry(future=self._executor.submit(work))
        log.info("已受理同步生成任务 %s（在途 %d）", task_id, len(self._entries))
        return task_id

    def fetch(self, task_id: str) -> GenerationResult | None:
        """查任务；未知 ID 返回 None。

        未知 ID 的两种成因**无法区分**：从未存在过，或进程重启后丢失。两者对调用
        方的处置相同（停止轮询、任务判失败），故不强行区分——区分不了的事不要假装
        区分。
        """
        with self._lock:
            entry = self._entries.get(task_id)
            if entry is not None and entry.finished_at is None and entry.future.done():
                # 首次观察到完成：记下时刻，TTL 由此起算。
                entry.finished_at = time.monotonic()
        if entry is None:
            return None

        if not entry.future.done():
            # 同步接口不提供进度，故只能报 0。客户端契约本就不承诺平滑进度
            # （ADR-023 §10），一个诚实的 0 不该被读成「卡住」。
            return GenerationResult(
                provider_request_id=task_id, state=TaskState.RUNNING
            )

        try:
            result = entry.future.result()
        except Exception as exc:  # noqa: BLE001 - 兜底转成领域失败
            return self._failure_result(task_id, exc)
        # 用真实的任务 ID 覆盖工作函数里可能留空的那个字段。
        return GenerationResult(
            provider_request_id=task_id,
            state=result.state,
            progress=result.progress,
            asset_url=result.asset_url,
            asset_url_expires_in_sec=result.asset_url_expires_in_sec,
            cover_url=result.cover_url,
            ratio=result.ratio,
            resolution=result.resolution,
            width_px=result.width_px,
            height_px=result.height_px,
            duration_ms=result.duration_ms,
            output_format=result.output_format,
            failure=result.failure,
            billing_tokens=result.billing_tokens,
        )

    @staticmethod
    def _failure_result(task_id: str, exc: BaseException) -> GenerationResult:
        """工作函数抛出的异常转成领域失败。

        已分类的服务商异常保留其分类；其余按可重试的上游问题处理，并把类型名带进
        `provider_code`——否则排查时只有一句空洞的错误消息。
        """
        from app.domain.generation import Retry, UserFacing

        from .client import OpenAiError

        if isinstance(exc, OpenAiError):
            failure = exc.failure
        else:
            log.exception("同步生成任务 %s 抛出未分类异常", task_id)
            failure = Failure(
                Retry.BACKOFF,
                UserFacing.UPSTREAM_UNAVAILABLE,
                type(exc).__name__,
                str(exc),
            )
        return GenerationResult(
            provider_request_id=task_id, state=TaskState.FAILED, failure=failure
        )

    def _evict_expired_locked(self) -> None:
        """清理已完成且超过 TTL 的条目。调用方必须已持锁。

        **只清已完成的**：在途任务无论多久都不清——它可能正在花钱，清掉会让它
        变成一个无人认领的开销。真正卡死的在途任务由线程池的上限兜住。
        """
        now = time.monotonic()
        # 顺带给「已完成但从未被 fetch 过」的条目补上观察时刻，否则它们的
        # finished_at 永远是 None，就永远不会过期——那正是内存泄漏。
        for entry in self._entries.values():
            if entry.finished_at is None and entry.future.done():
                entry.finished_at = now
        expired = [
            task_id
            for task_id, entry in self._entries.items()
            if entry.finished_at is not None
            and now - entry.finished_at > self._completed_ttl_sec
        ]
        for task_id in expired:
            del self._entries[task_id]
        if expired:
            log.debug("清理了 %d 个已完成的同步任务条目", len(expired))

    @property
    def size(self) -> int:
        with self._lock:
            return len(self._entries)

    def shutdown(self) -> None:
        """停止接受新任务。

        `wait=False`：在途的同步调用可能还要跑几分钟，进程退出时不等它们——结果
        反正已经无处可交付了。它们会随进程一起消失，这正是上面第 1 点登记的风险。
        """
        self._executor.shutdown(wait=False, cancel_futures=True)