from __future__ import annotations import queue import threading import uuid from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from .schemas import JobDetail, JobEvent, JobStatus @dataclass class JobState: status: JobStatus events: list[JobEvent] = field(default_factory=list) subscribers: list[queue.Queue] = field(default_factory=list) class JobManager: def __init__(self) -> None: self._lock = threading.Lock() self._jobs: dict[str, JobState] = {} def create_job(self) -> str: job_id = uuid.uuid4().hex now = datetime.now(timezone.utc) status = JobStatus( job_id=job_id, status="queued", stage="queued", progress_percent=0, message="任务已创建", created_at=now, updated_at=now, artifacts={}, error="", ) with self._lock: self._jobs[job_id] = JobState(status=status) return job_id def exists(self, job_id: str) -> bool: with self._lock: return job_id in self._jobs def get_status(self, job_id: str) -> JobStatus: with self._lock: return self._jobs[job_id].status def get_detail(self, job_id: str) -> JobDetail: with self._lock: state = self._jobs[job_id] return JobDetail(status=state.status, recent_events=state.events[-100:]) def set_artifacts(self, job_id: str, artifacts: dict[str, str]) -> None: with self._lock: status = self._jobs[job_id].status status.artifacts = artifacts status.updated_at = datetime.now(timezone.utc) def set_status( self, job_id: str, *, status: str, stage: str, progress_percent: int, message: str, error: str | None = None, ) -> None: with self._lock: s = self._jobs[job_id].status s.status = status s.stage = stage s.progress_percent = progress_percent s.message = message s.updated_at = datetime.now(timezone.utc) if error is not None: s.error = error def add_event(self, job_id: str, *, kind: str, stage: str, progress_percent: int, message: str) -> None: event = JobEvent( type=kind, stage=stage, progress_percent=progress_percent, message=message, timestamp=datetime.now(timezone.utc), ) with self._lock: state = self._jobs[job_id] state.events.append(event) state.status.stage = stage state.status.progress_percent = progress_percent state.status.message = message state.status.updated_at = event.timestamp for sub in state.subscribers: sub.put(event) def subscribe(self, job_id: str) -> queue.Queue: q: queue.Queue = queue.Queue() with self._lock: self._jobs[job_id].subscribers.append(q) return q def unsubscribe(self, job_id: str, q: queue.Queue) -> None: with self._lock: subs = self._jobs[job_id].subscribers if q in subs: subs.remove(q) @staticmethod def resolve_artifact_path(status: JobStatus, kind: str) -> Path: if kind not in status.artifacts: raise KeyError(kind) p = status.artifacts[kind] if not p: raise FileNotFoundError(kind) return Path(p)