329 lines
14 KiB
Python
329 lines
14 KiB
Python
"""异步后台任务注册表（T1：后端异步化核心）。
|
||
|
||
设计原则（对齐 v3 方案 D1-D6）：
|
||
- asyncio 原生，无 Celery/Redis/外部队列
|
||
- 每个 request_id -> TaskInfo（状态/开始时间/结果或错误）
|
||
- 任务写 runs/{id}/workspace.json，SSE 生成器只读该文件（D3：不改 pipeline）
|
||
- 定期清理已完成任务（防内存泄漏）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any, AsyncGenerator, Dict, Optional
|
||
|
||
# runs/ 目录（与 pipeline.py 默认一致）
|
||
RUNS_DIR = Path("runs")
|
||
|
||
# 任务状态
|
||
STATE_PENDING = "pending"
|
||
STATE_RUNNING = "running"
|
||
STATE_DONE = "done"
|
||
STATE_FAILED = "failed"
|
||
|
||
|
||
@dataclass
|
||
class TaskInfo:
|
||
"""一个后台任务的状态快照。"""
|
||
request_id: str
|
||
state: str = STATE_PENDING # pending | running | done | failed
|
||
started_at: float = 0.0 # time.time()
|
||
finished_at: float = 0.0 # time.time()（done/failed 时）
|
||
error: Optional[str] = None # failed 时错误信息
|
||
# PipelineResult 字段（done 时填充）
|
||
response: Optional[str] = None
|
||
status: Optional[str] = None # done | fast_path | escalated | failed
|
||
fast_path: bool = False
|
||
rounds_used: int = 0
|
||
api_input_tokens: int = 0
|
||
api_output_tokens: int = 0
|
||
cost_est: float = 0.0
|
||
model_used: Optional[str] = None
|
||
latency_ms: float = 0.0
|
||
workspace_path: Optional[str] = None
|
||
error_detail: Optional[str] = None
|
||
route: list = field(default_factory=list)
|
||
|
||
def to_dict(self) -> Dict[str, Any]:
|
||
return {
|
||
"request_id": self.request_id,
|
||
"state": self.state,
|
||
"started_at": self.started_at,
|
||
"finished_at": self.finished_at,
|
||
"error": self.error,
|
||
"response": self.response,
|
||
"status": self.status,
|
||
"fast_path": self.fast_path,
|
||
"rounds_used": self.rounds_used,
|
||
"api_input_tokens": self.api_input_tokens,
|
||
"api_output_tokens": self.api_output_tokens,
|
||
"cost_est": self.cost_est,
|
||
"model_used": self.model_used,
|
||
"latency_ms": round(self.latency_ms, 2),
|
||
"workspace_path": self.workspace_path,
|
||
"route": self.route,
|
||
"error_detail": self.error_detail,
|
||
}
|
||
|
||
def to_result_event(self) -> Dict[str, Any]:
|
||
"""终态 SSE result 事件 payload。"""
|
||
return {
|
||
"type": "result",
|
||
"request_id": self.request_id,
|
||
"response": self.response or "",
|
||
"status": self.status,
|
||
"fast_path": self.fast_path,
|
||
"rounds_used": self.rounds_used,
|
||
"api_input_tokens": self.api_input_tokens,
|
||
"api_output_tokens": self.api_output_tokens,
|
||
"cost_est": self.cost_est,
|
||
"model_used": self.model_used,
|
||
"latency_ms": round(self.latency_ms, 2),
|
||
"route": self.route,
|
||
"error": self.error,
|
||
}
|
||
|
||
|
||
class JobStore:
|
||
"""asyncio 后台任务注册表。
|
||
|
||
线程安全（asyncio 事件循环单线程，不需要额外锁）。
|
||
任务由 register() 注册、由 _task_done() 填充结果。
|
||
SSE 生成器调用 watch_file() 轮询 workspace.json 变化。
|
||
"""
|
||
|
||
def __init__(self, max_age_seconds: float = 3600.0, max_running: int = 50):
|
||
"""
|
||
Args:
|
||
max_age_seconds: 已完成任务在内存中保留时间（秒），超时后自动清理
|
||
max_running: 最大同时运行任务数，超出后拒绝新任务
|
||
"""
|
||
self._tasks: Dict[str, TaskInfo] = {}
|
||
self._asyncio_tasks: Dict[str, "asyncio.Task[None]"] = {} # request_id -> asyncio.Task
|
||
self._cancel_events: Dict[str, asyncio.Event] = {} # request_id -> cancel Event
|
||
self.max_age_seconds = max_age_seconds
|
||
self.max_running = max_running
|
||
self._poll_interval = 0.3 # workspace.json 轮询间隔（秒）
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# 公共 API
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
def register(self, request_id: str) -> tuple[bool, str]:
|
||
"""注册一个 pending 任务。返回 (True, "") 成功，(False, reason) 容量满。"""
|
||
running = [t for t in self._tasks.values() if t.state == STATE_RUNNING]
|
||
if len(running) >= self.max_running:
|
||
return False, f"同时运行任务已达上限 {self.max_running}，请稍后重试"
|
||
if request_id in self._tasks:
|
||
return False, f"任务 {request_id} 已存在"
|
||
self._tasks[request_id] = TaskInfo(
|
||
request_id=request_id,
|
||
state=STATE_PENDING,
|
||
started_at=time.time(),
|
||
)
|
||
return True, ""
|
||
|
||
def get(self, request_id: str) -> Optional[TaskInfo]:
|
||
"""查询任务状态。"""
|
||
return self._tasks.get(request_id)
|
||
|
||
def list_all(self) -> Dict[str, TaskInfo]:
|
||
"""列出所有任务（含已完成）。"""
|
||
return dict(self._tasks)
|
||
|
||
def list_running(self) -> Dict[str, TaskInfo]:
|
||
"""只列出运行中任务。"""
|
||
return {k: v for k, v in self._tasks.items() if v.state == STATE_RUNNING}
|
||
|
||
def submit(
|
||
self,
|
||
request_id: str,
|
||
coro, # type: asyncio.coroutine
|
||
) -> "asyncio.Task[None]":
|
||
"""提交协程到后台运行；内部注册 asyncio.Task。"""
|
||
task = asyncio.create_task(self._run_wrapper(request_id, coro))
|
||
self._asyncio_tasks[request_id] = task
|
||
return task
|
||
|
||
def new_cancel(self, request_id: str) -> asyncio.Event:
|
||
"""为 SSE 连接创建一个取消事件。"""
|
||
evt = asyncio.Event()
|
||
self._cancel_events[request_id] = evt
|
||
return evt
|
||
|
||
def get_cancel(self, request_id: str) -> Optional[asyncio.Event]:
|
||
"""获取已有取消事件。"""
|
||
return self._cancel_events.get(request_id)
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# 内部：任务运行包装
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
async def _run_wrapper(self, request_id: str, coro):
|
||
"""把用户协程包装成可追踪的任务：更新状态、捕获异常、清理。"""
|
||
info = self._tasks.get(request_id)
|
||
if info is None:
|
||
return
|
||
info.state = STATE_RUNNING
|
||
try:
|
||
await coro
|
||
except Exception as exc: # pragma: no cover
|
||
if info:
|
||
info.state = STATE_FAILED
|
||
info.finished_at = time.time()
|
||
info.error = f"unhandled:{type(exc).__name__}:{exc}"
|
||
finally:
|
||
# 清理 asyncio task 引用
|
||
self._asyncio_tasks.pop(request_id, None)
|
||
# 定期 GC 已完成任务
|
||
self._cleanup_aged()
|
||
|
||
def _task_done(
|
||
self,
|
||
request_id: str,
|
||
result: "PipelineResult", # from pipeline.PipelineResult
|
||
) -> None:
|
||
"""任务正常完成时由调用方调用，写入结果。"""
|
||
info = self._tasks.get(request_id)
|
||
if info is None:
|
||
return
|
||
info.state = STATE_DONE
|
||
info.finished_at = time.time()
|
||
info.response = result.response
|
||
info.status = result.status
|
||
info.fast_path = result.fast_path
|
||
info.rounds_used = result.rounds_used
|
||
info.api_input_tokens = result.api_input_tokens
|
||
info.api_output_tokens = result.api_output_tokens
|
||
info.cost_est = result.cost_est
|
||
info.model_used = result.model_used
|
||
info.latency_ms = result.latency_ms
|
||
info.workspace_path = result.workspace_path
|
||
info.route = list(result.route) if result.route else []
|
||
if result.error:
|
||
info.error_detail = result.error
|
||
|
||
def _task_failed(
|
||
self,
|
||
request_id: str,
|
||
error: str,
|
||
detail: Optional[str] = None,
|
||
) -> None:
|
||
"""任务异常结束时调用。"""
|
||
info = self._tasks.get(request_id)
|
||
if info is None:
|
||
return
|
||
info.state = STATE_FAILED
|
||
info.finished_at = time.time()
|
||
info.error = error
|
||
info.error_detail = detail
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# SSE 专用：workspace.json 文件轮询
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
async def watch_workspace(
|
||
self,
|
||
request_id: str,
|
||
cancel_event: asyncio.Event,
|
||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||
"""监视 runs/{id}/workspace.json，yield 单条事件。
|
||
|
||
事件类型：
|
||
- {"type": "status", "value": <state>, "request_id": <id>}
|
||
- {"type": "workspace", "version": <n>, "workspace": <dict>}
|
||
- {"type": "error", "detail": <str>}
|
||
|
||
当任务进入 done/failed 状态或 cancel_event 被 set 时停止。
|
||
"""
|
||
ws_path = RUNS_DIR / request_id / "workspace.json"
|
||
seen_mtime: float = 0.0
|
||
seen_size: int = 0
|
||
|
||
while not cancel_event.is_set():
|
||
info = self.get(request_id)
|
||
# 检查终态
|
||
if info and info.state in (STATE_DONE, STATE_FAILED):
|
||
if info.state == STATE_FAILED:
|
||
yield {"type": "error", "detail": info.error or "任务失败"}
|
||
# result 由 /runs/{id}/status 提供，此处只推送终态 status
|
||
yield {"type": "status", "value": info.state, "request_id": request_id}
|
||
break
|
||
|
||
# 读文件变化（mtime + size 双检）
|
||
if ws_path.exists():
|
||
try:
|
||
stat = ws_path.stat()
|
||
if stat.st_mtime != seen_mtime or stat.st_size != seen_size:
|
||
raw = ws_path.read_text(encoding="utf-8")
|
||
data = json.loads(raw)
|
||
version = int(data.get("meta", {}).get("round", 0))
|
||
seen_mtime = stat.st_mtime
|
||
seen_size = stat.st_size
|
||
yield {
|
||
"type": "workspace",
|
||
"version": version,
|
||
"state": info.state if info else STATE_RUNNING,
|
||
"request_id": request_id,
|
||
"workspace": data,
|
||
}
|
||
except (json.JSONDecodeError, OSError):
|
||
pass # 文件正在写入，忽略
|
||
|
||
await asyncio.sleep(self._poll_interval)
|
||
|
||
# 连接关闭前最后推一次终态
|
||
info = self.get(request_id)
|
||
if info:
|
||
yield {"type": "status", "value": info.state, "request_id": request_id}
|
||
|
||
# ─────────────────────────────────────────────────────────────────
|
||
# 内存清理
|
||
# ─────────────────────────────────────────────────────────────────
|
||
|
||
def _cleanup_aged(self) -> None:
|
||
"""删除超过 max_age_seconds 的已完成任务引用。"""
|
||
now = time.time()
|
||
to_remove = [
|
||
rid for rid, info in self._tasks.items()
|
||
if info.state in (STATE_DONE, STATE_FAILED)
|
||
and (now - info.finished_at) > self.max_age_seconds
|
||
]
|
||
for rid in to_remove:
|
||
self._tasks.pop(rid, None)
|
||
|
||
def cancel(self, request_id: str) -> bool:
|
||
"""取消运行中的任务。返回 True 找到并取消，False 未找到。"""
|
||
task = self._asyncio_tasks.get(request_id)
|
||
if task is None:
|
||
return False
|
||
task.cancel()
|
||
info = self.get(request_id)
|
||
if info:
|
||
info.state = STATE_FAILED
|
||
info.finished_at = time.time()
|
||
info.error = "cancelled_by_user"
|
||
return True
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# 全局单例（gateway 进程内共享）
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
_store: Optional[JobStore] = None
|
||
|
||
|
||
def get_job_store() -> JobStore:
|
||
global _store
|
||
if _store is None:
|
||
_store = JobStore(max_age_seconds=3600.0, max_running=50)
|
||
return _store
|
||
|
||
|
||
def reset_job_store() -> None:
|
||
"""测试用：重置全局注册表。"""
|
||
global _store
|
||
_store = None
|