8ecc701d5e
实现后台任务调度器(scheduler.py)和任务运行器(task_runner.py)，支持长时间运行任务的异步执行和状态跟踪 新增多种工具支持：Shell命令执行、文件操作(读写/搜索/发送)、网页搜索/问答、定时提醒等 扩展README和ROADMAP文档，描述新功能和未来多主机架构规划 在配置文件中添加METASO_API_KEY支持秘塔AI搜索功能 优化代理逻辑，自动识别通用问题直接回答而不创建会话
148 lines
4.3 KiB
Python
148 lines
4.3 KiB
Python
"""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
|
|
|
|
@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()
|
|
|
|
async def submit(
|
|
self,
|
|
coro: Callable[[], Any],
|
|
description: str,
|
|
notify_chat_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,
|
|
)
|
|
|
|
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:
|
|
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()
|