"""Task state management. Tracks long-running tasks (e.g. graph build) so callers can poll progress. """ import uuid import threading from datetime import datetime from enum import Enum from typing import Dict, Any, Optional from dataclasses import dataclass, field from ..utils.locale import t class TaskStatus(str, Enum): """Task status enum.""" PENDING = "pending" # waiting PROCESSING = "processing" # in progress COMPLETED = "completed" # finished successfully FAILED = "failed" # finished with error @dataclass class Task: """Task data class.""" task_id: str task_type: str status: TaskStatus created_at: datetime updated_at: datetime progress: int = 0 # overall progress percentage 0-100 message: str = "" # human-readable status message result: Optional[Dict] = None # task result payload error: Optional[str] = None # error message when failed metadata: Dict = field(default_factory=dict) # arbitrary caller metadata progress_detail: Dict = field(default_factory=dict) # fine-grained progress info def to_dict(self) -> Dict[str, Any]: """Serialize the task to a JSON-friendly dict.""" return { "task_id": self.task_id, "task_type": self.task_type, "status": self.status.value, "created_at": self.created_at.isoformat(), "updated_at": self.updated_at.isoformat(), "progress": self.progress, "message": self.message, "progress_detail": self.progress_detail, "result": self.result, "error": self.error, "metadata": self.metadata, } class TaskManager: """Thread-safe singleton task registry.""" _instance = None _lock = threading.Lock() def __new__(cls): if cls._instance is None: with cls._lock: if cls._instance is None: cls._instance = super().__new__(cls) cls._instance._tasks: Dict[str, Task] = {} cls._instance._task_lock = threading.Lock() return cls._instance def create_task(self, task_type: str, metadata: Optional[Dict] = None) -> str: """Create a new task. Args: task_type: Task type identifier. metadata: Optional caller-supplied metadata. Returns: The newly created task id. """ task_id = str(uuid.uuid4()) now = datetime.now() task = Task( task_id=task_id, task_type=task_type, status=TaskStatus.PENDING, created_at=now, updated_at=now, metadata=metadata or {} ) with self._task_lock: self._tasks[task_id] = task return task_id def get_task(self, task_id: str) -> Optional[Task]: """Return the task for ``task_id`` or ``None`` if unknown.""" with self._task_lock: return self._tasks.get(task_id) def update_task( self, task_id: str, status: Optional[TaskStatus] = None, progress: Optional[int] = None, message: Optional[str] = None, result: Optional[Dict] = None, error: Optional[str] = None, progress_detail: Optional[Dict] = None ): """Update mutable fields on an existing task. Args: task_id: Task id to update. status: New status, if changing. progress: New overall progress (0-100), if changing. message: New status message, if changing. result: New result payload, if changing. error: New error message, if changing. progress_detail: New fine-grained progress info, if changing. """ with self._task_lock: task = self._tasks.get(task_id) if task: task.updated_at = datetime.now() if status is not None: task.status = status if progress is not None: task.progress = progress if message is not None: task.message = message if result is not None: task.result = result if error is not None: task.error = error if progress_detail is not None: task.progress_detail = progress_detail def complete_task(self, task_id: str, result: Dict): """Mark a task as completed and attach the result.""" self.update_task( task_id, status=TaskStatus.COMPLETED, progress=100, message=t('progress.taskComplete'), result=result ) def fail_task(self, task_id: str, error: str): """Mark a task as failed and attach the error message.""" self.update_task( task_id, status=TaskStatus.FAILED, message=t('progress.taskFailed'), error=error ) def list_tasks(self, task_type: Optional[str] = None) -> list: """List tasks, optionally filtered by ``task_type``, newest first.""" with self._task_lock: tasks = list(self._tasks.values()) if task_type: tasks = [t for t in tasks if t.task_type == task_type] return [t.to_dict() for t in sorted(tasks, key=lambda x: x.created_at, reverse=True)] def cleanup_old_tasks(self, max_age_hours: int = 24): """Drop completed/failed tasks older than ``max_age_hours``.""" from datetime import timedelta cutoff = datetime.now() - timedelta(hours=max_age_hours) with self._task_lock: old_ids = [ tid for tid, task in self._tasks.items() if task.created_at < cutoff and task.status in [TaskStatus.COMPLETED, TaskStatus.FAILED] ] for tid in old_ids: del self._tasks[tid]