"""Background task runner for long-running operations.""" from __future__ import annotations import asyncio import logging import time import uuid from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable, Dict, Optional logger = logging.getLogger(__name__) class TaskStatus(str, Enum): PENDING = "pending" RUNNING = "running" COMPLETED = "completed" FAILED = "failed" @dataclass class BackgroundTask: task_id: str description: str started_at: float status: TaskStatus = TaskStatus.PENDING completed_at: Optional[float] = None result: Optional[str] = None error: Optional[str] = None notify_chat_id: Optional[str] = None user_id: Optional[str] = None @property def elapsed(self) -> float: if self.completed_at: return self.completed_at - self.started_at return time.time() - self.started_at class TaskRunner: """Singleton that manages background tasks with Feishu notifications.""" def __init__(self) -> None: self._tasks: Dict[str, BackgroundTask] = {} self._lock = asyncio.Lock() self._notification_handler: Optional[Callable] = None def set_notification_handler(self, handler: Optional[Callable]) -> None: """Set custom notification handler for M3 mode (host client -> router).""" self._notification_handler = handler async def submit( self, coro: Callable[[], Any], description: str, notify_chat_id: Optional[str] = None, user_id: Optional[str] = None, ) -> str: """Submit a coroutine as a background task.""" task_id = str(uuid.uuid4())[:8] task = BackgroundTask( task_id=task_id, description=description, started_at=time.time(), status=TaskStatus.PENDING, notify_chat_id=notify_chat_id, user_id=user_id, ) async with self._lock: self._tasks[task_id] = task asyncio.create_task(self._run_task(task_id, coro)) logger.info("Submitted background task %s: %s", task_id, description) return task_id async def _run_task(self, task_id: str, coro: Callable[[], Any]) -> None: """Execute a task and send notification on completion.""" async with self._lock: task = self._tasks.get(task_id) if not task: return task.status = TaskStatus.RUNNING try: result = await coro async with self._lock: task.status = TaskStatus.COMPLETED task.completed_at = time.time() task.result = str(result)[:2000] if result else None logger.info("Task %s completed in %.1fs", task_id, task.elapsed) except Exception as exc: async with self._lock: task.status = TaskStatus.FAILED task.completed_at = time.time() task.error = str(exc)[:500] logger.exception("Task %s failed: %s", task_id, exc) if task.notify_chat_id: if self._notification_handler: await self._notification_handler(task) else: await self._send_notification(task) async def _send_notification(self, task: BackgroundTask) -> None: """Send Feishu notification about task completion.""" from bot.feishu import send_text if task.status == TaskStatus.COMPLETED: emoji = "✅" status_text = "done" else: emoji = "❌" status_text = "failed" elapsed = int(task.elapsed) msg = f"{emoji} Task #{task.task_id} {status_text} ({elapsed}s)\n{task.description}" if task.result: truncated = task.result[:800] if len(task.result) > 800: truncated += "..." msg += f"\n\n```\n{truncated}\n```" elif task.error: msg += f"\n\n**Error:** {task.error}" try: await send_text(task.notify_chat_id, "chat_id", msg) except Exception: logger.exception("Failed to send notification for task %s", task.task_id) def get_task(self, task_id: str) -> Optional[BackgroundTask]: return self._tasks.get(task_id) def list_tasks(self, limit: int = 20) -> list[dict]: tasks = sorted( self._tasks.values(), key=lambda t: t.started_at, reverse=True, )[:limit] return [ { "task_id": t.task_id, "description": t.description, "status": t.status.value, "elapsed": int(t.elapsed), "started_at": t.started_at, } for t in tasks ] task_runner = TaskRunner()