aps-agent/server/agent_core/async_jobs.py

# ============================================================
# 轻量异步任务队列（moduleId: core-async-jobs, 可重生 ✅）
# 矩阵 88 行剩余：敏感性 / 蒙特卡洛重算放后台（提交 / 轮询 / 取消 / 结果）
# 线程模型：daemon worker 线程执行；fn 签名 fn(record)，可检查
#   record.cancel_requested 提前退出（或抛 JobCancelled）实现可取消长任务。
# 与网关解耦：本模块只做队列与状态机，业务闭包由调用方注入。
# ============================================================
from __future__ import annotations

import itertools
import threading
import time
import uuid
from collections.abc import Callable
from datetime import UTC, datetime
from typing import Any, Literal

from pydantic import BaseModel, Field

from server.timeutil import fmt_dt

JobStatus = Literal["pending", "running", "done", "failed", "cancelled"]
_TERMINAL = ("done", "failed", "cancelled")


class JobCancelled(Exception):
    """任务主动取消：worker 捕获后状态归为 cancelled。"""


class JobQueueFullError(RuntimeError):
    """在途任务达到上限，拒绝继续创建线程。"""


def _now() -> str:
    # DTZ005：统一取本地时区的 aware 时间，墙钟与仓库其余模块一致
    return fmt_dt(datetime.now(UTC).astimezone())


class JobRecord(BaseModel):
    """任务记录：状态机 pending → running → done/failed/cancelled。"""

    job_id: str
    kind: str
    status: JobStatus = "pending"
    progress: int = 0
    params: dict[str, Any] = Field(default_factory=dict)
    actor: str = "planner"
    result: Any = None
    error: str | None = None
    created_at: str = ""
    started_at: str | None = None
    finished_at: str | None = None
    cancel_requested: bool = False
    tenant_uuid: str = "platform"
    project_id: str = "default"

    def raise_if_cancelled(self) -> None:
        if self.cancel_requested:
            raise JobCancelled()


class JobQueue:
    """轻量异步任务队列：提交 / 轮询 / 取消 / 结果 / 统计。"""

    def __init__(self, *, max_finished: int = 200, max_concurrent: int = 2,
                 max_inflight: int = 8) -> None:
        self._lock = threading.RLock()
        self._jobs: dict[str, JobRecord] = {}
        self._seq = itertools.count(1)
        self.max_finished = max_finished
        self.max_concurrent = max(1, int(max_concurrent))
        self.max_inflight = max(self.max_concurrent, int(max_inflight))
        self._slots = threading.BoundedSemaphore(self.max_concurrent)
        self._inflight_threads = 0

    # ---------------- 提交 ----------------
    def submit(
        self,
        kind: str,
        fn: Callable[[JobRecord], Any],
        *,
        params: dict[str, Any] | None = None,
        actor: str = "planner",
        tenant_uuid: str = "platform",
        project_id: str = "default",
    ) -> str:
        """提交任务：立即起 daemon 线程执行，返回 job_id。"""
        job_id = f"job-{next(self._seq):04d}-{uuid.uuid4().hex[:6]}"
        with self._lock:
            if self._inflight_threads >= self.max_inflight:
                raise JobQueueFullError(f"异步任务在途已达上限 {self.max_inflight}")
            record = JobRecord(
                job_id=job_id, kind=kind, params=params or {}, actor=actor,
                tenant_uuid=tenant_uuid, project_id=project_id, created_at=_now(),
            )
            self._jobs[job_id] = record
            self._inflight_threads += 1
        thread = threading.Thread(
            target=self._worker, args=(job_id, fn),
            name=f"aps-job-{job_id}", daemon=True,
        )
        try:
            thread.start()
        except BaseException:
            with self._lock:
                self._jobs.pop(job_id, None)
                self._inflight_threads = max(0, self._inflight_threads - 1)
            raise
        return job_id

    def _worker(self, job_id: str, fn: Callable[[JobRecord], Any]) -> None:
        try:
            with self._slots:
                with self._lock:
                    record = self._jobs[job_id]
                    if record.status == "cancelled" or record.cancel_requested:
                        self._finish(job_id, "cancelled")
                        return
                    record.status = "running"; record.started_at = _now()
                try:
                    result = fn(record)
                    with self._lock:
                        current = self._jobs[job_id]
                        if current.cancel_requested:
                            self._finish(job_id, "cancelled")
                        else:
                            current.result = result; current.progress = 100
                            self._finish(job_id, "done")
                except JobCancelled:
                    self._finish(job_id, "cancelled")
                except Exception as exc:  # noqa: BLE001
                    with self._lock:
                        current = self._jobs[job_id]
                        if current.cancel_requested:
                            self._finish(job_id, "cancelled")
                        else:
                            current.error = str(exc)
                            self._finish(job_id, "failed")
        finally:
            with self._lock:
                self._inflight_threads = max(0, self._inflight_threads - 1)
            self._prune()

    def _finish(self, job_id: str, status: JobStatus) -> None:
        with self._lock:
            record = self._jobs[job_id]
            record.status = status
            record.finished_at = _now()
            if status == "cancelled":
                record.result = None
                record.error = "cancelled"

    # ---------------- 轮询 / 结果 ----------------
    def poll(self, job_id: str, *, tenant_uuid: str | None = None,
             project_id: str | None = None,
             actor: str | None = None) -> dict[str, Any] | None:
        """查询任务快照（含 result/error）；不存在返回 None。"""
        with self._lock:
            record = self._jobs.get(job_id)
            if record and tenant_uuid is not None and record.tenant_uuid != tenant_uuid: return None
            if record and project_id is not None and record.project_id != project_id: return None
            if record and actor is not None and record.actor != actor: return None
            return record.model_dump(mode="json") if record else None

    def wait(self, job_id: str, timeout: float = 15.0) -> dict[str, Any]:
        """阻塞轮询至终态（done/failed/cancelled）。"""
        deadline = time.time() + timeout
        while True:
            record = self.poll(job_id)
            if record is None:
                raise KeyError(job_id)
            if record["status"] in _TERMINAL:
                return record
            if time.time() >= deadline:
                raise TimeoutError(f"任务 {job_id} 等待超时（{timeout}s）")
            time.sleep(0.05)

    # ---------------- 取消 ----------------
    def cancel(self, job_id: str, *, tenant_uuid: str | None = None,
               project_id: str | None = None,
               actor: str | None = None) -> bool:
        """取消任务：pending 立即取消；running 置取消标记（worker 收尾时落 cancelled）；
        已终态或不存在返回 False。"""
        with self._lock:
            record = self._jobs.get(job_id)
            if record is None or record.status in _TERMINAL:
                return False
            if tenant_uuid is not None and record.tenant_uuid != tenant_uuid: return False
            if project_id is not None and record.project_id != project_id: return False
            if actor is not None and record.actor != actor: return False
            record.cancel_requested = True
            if record.status == "pending":
                record.status = "cancelled"
                record.finished_at = _now()
            return True

    # ---------------- 列表 / 统计 ----------------
    def list(self, status: str | None = None, *, tenant_uuid: str | None = None,
             project_id: str | None = None,
             actor: str | None = None) -> list[dict[str, Any]]:
        with self._lock:
            items = [r.model_dump(mode="json") for r in self._jobs.values()]
        if status:
            items = [i for i in items if i["status"] == status]
        if tenant_uuid is not None: items = [i for i in items if i["tenant_uuid"] == tenant_uuid]
        if project_id is not None: items = [i for i in items if i["project_id"] == project_id]
        if actor is not None: items = [i for i in items if i["actor"] == actor]
        items.sort(key=lambda i: i["created_at"] or "", reverse=True)
        return items

    def stats(self, *, tenant_uuid: str | None = None,
              project_id: str | None = None,
              actor: str | None = None) -> dict[str, int]:
        with self._lock:
            counts: dict[str, int] = {
                "total": len(self._jobs), "pending": 0, "running": 0,
                "done": 0, "failed": 0, "cancelled": 0,
            }
            for r in self._jobs.values():
                if tenant_uuid is not None and r.tenant_uuid != tenant_uuid: continue
                if project_id is not None and r.project_id != project_id: continue
                if actor is not None and r.actor != actor: continue
                counts[r.status] = counts.get(r.status, 0) + 1
            counts["total"] = sum(counts[s] for s in ("pending", "running", "done", "failed", "cancelled"))
        return counts

    def _prune(self) -> None:
        """只保留最多 max_finished 条终态任务（按完成时间淘汰最旧）。"""
        with self._lock:
            finished = sorted(
                (r for r in self._jobs.values() if r.status in _TERMINAL),
                key=lambda r: r.finished_at or "",
            )
            overflow = len(finished) - self.max_finished
            for r in finished[:overflow]:
                self._jobs.pop(r.job_id, None)

    def clear(self) -> None:
        with self._lock:
            self._jobs.clear()


_queue: JobQueue | None = None


def get_job_queue() -> JobQueue:
    global _queue
    if _queue is None:
        _queue = JobQueue()
    return _queue


def reset_job_queue() -> None:
    """测试用：丢弃单例，下次重建。"""
    global _queue
    _queue = None