"""Background task manager for async task tracking.""" import asyncio import logging from dataclasses import dataclass, field from typing import Dict, Coroutine from core.constants import TaskStatus logger = logging.getLogger(__name__) @dataclass class TaskInfo: """Information about a running task.""" task_id: str task_type: str coroutine: Coroutine = None status: TaskStatus = TaskStatus.RUNNING created_at: float = field(default_factory=lambda: asyncio.get_running_loop().time()) metadata: dict = field(default_factory=dict) class BackgroundTaskManager: """ Manages background async tasks with tracking and cancellation support. Features: - Start and track async tasks - Cancel individual or all tasks - Get task status and metadata - Prevent duplicate tasks """ def __init__(self): self._tasks: Dict[str, asyncio.Task] = {} self._task_info: Dict[str, TaskInfo] = {} self._lock = asyncio.Lock() async def start_task( self, task_id: str, coro: Coroutine, task_type: str = "unknown", metadata: dict = None ) -> bool: """ Start a background task. Args: task_id: Unique task identifier coro: Coroutine to run task_type: Type of task metadata: Additional task metadata Returns: True if started successfully, False if already running """ async with self._lock: if task_id in self._tasks and not self._tasks[task_id].done(): logger.warning(f"Task {task_id} is already running") return False task = asyncio.create_task(coro) self._tasks[task_id] = task self._task_info[task_id] = TaskInfo( task_id=task_id, task_type=task_type, coroutine=coro, metadata=metadata or {} ) task.add_done_callback(lambda t, tid=task_id: self._on_task_done(tid, t)) logger.info(f"Started background task: {task_id} ({task_type})") return True def _on_task_done(self, task_id: str, task: asyncio.Task) -> None: if task.cancelled(): logger.info(f"Task {task_id} cancelled") return exc = task.exception() if exc: logger.error(f"Task {task_id} failed: {exc}") else: logger.info(f"Task {task_id} completed successfully") async def cancel_task(self, task_id: str) -> bool: """ Cancel a specific task. Args: task_id: Task identifier to cancel Returns: True if task was cancelled, False if not found """ async with self._lock: task = self._tasks.get(task_id) if task and not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass logger.info(f"Task cancelled: {task_id}") return True return False async def cancel_all(self) -> int: """ Cancel all running tasks. Returns: Number of tasks cancelled """ cancelled = 0 async with self._lock: for task in list(self._tasks.values()): if not task.done(): task.cancel() try: await task except asyncio.CancelledError: pass cancelled += 1 self._tasks.clear() self._task_info.clear() logger.info(f"All tasks cancelled: {cancelled}") return cancelled def get_active_tasks(self) -> Dict[str, dict]: """ Get all active tasks with their metadata. Returns: Dictionary of active tasks """ active = {} for task_id, task in self._tasks.items(): if not task.done(): info = self._task_info.get(task_id) active[task_id] = { "info": info.metadata if info else {}, "status": TaskStatus.RUNNING.value, "type": info.task_type if info else "unknown", "created_at": info.created_at if info else 0 } return active @property def active_count(self) -> int: """Number of currently active tasks.""" return sum(1 for t in self._tasks.values() if not t.done()) def has_active_tasks(self) -> bool: """Check if there are any active tasks.""" return self.active_count > 0