feat(core): 定时任务薄封装 scheduler / 跨插件任务槽 task_slots / 出站媒体显式发送方式
- scheduler.py:apscheduler 薄封装(统一 Asia/Shanghai、misfire_grace_time、起止/异常日志) - task_slots.py:长任务任务槽(进度上报 / 同槽去重 / SSE 推送 / 取消), video_analysis temp 清理等长任务接入 - outbound_media:新增 HEXI_MEDIA_MODE 显式发送方式(http/base64/file,/hub 可改), 修复 alconna 对 base64 载荷补 https:// 前缀导致的 fetch failed(repair_mangled_url) - learning_chat 定时任务迁入 hexi.core.scheduler - 测试:新增 test_scheduler / test_task_slots,扩充 test_outbound_media Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
1 parent
a789a44f5d
commit
0ac1958932
8 files changed
+1353
-12
No files matched your search
@@ -9,6 +9,7 @@ from . import ( # noqa: E402,F401
|
||||
plugin_control,
|
||||
plugin_manager,
|
||||
rate_limit,
|
||||
task_slots,
|
||||
)
|
||||
|
||||
|
||||
|
||||
+112
-10
@@ -10,6 +10,10 @@
|
||||
`bot.send` / alconna `finish` / 直接 `send_group_msg` 全从这里过)上,
|
||||
在真正调用适配器前改写发送类 API 里的本地文件:
|
||||
|
||||
发送方式由 `HEXI_MEDIA_MODE` 显式配置(`http` / `base64` / `file`,见 `_mode()`,
|
||||
/hub 全局设置页可改):默认 `http`,即下面的两级改写;`base64` 跳过第 1 级;
|
||||
`file` 整体不改写(协议端就在本机时才用)。
|
||||
|
||||
1. **默认改写为 `http://<本机可达IP>:<端口>/media/<token>`**,由协议端主动来拉。
|
||||
地址取自 `driver.config` 的 host/port + 本机网卡,与 /hub 首页「协议端配对」
|
||||
展示的是同一来源;反向 WS 部署下协议端本就连得到这个地址(它连的就是它),
|
||||
@@ -29,6 +33,14 @@
|
||||
且**只做 URL 不做 base64**——文件体积不适合内联。覆盖这条让
|
||||
`video_analysis` 的群文件「本地直传」兜底和 alconna 的 `$onebot11:file`
|
||||
一起受益(后者传的是 `Path.as_posix()` 裸路径)。
|
||||
|
||||
另有一类「必死值」需要顺手还原:alconna 的 `Media.__post_init__` 会给解析不出
|
||||
hostname 的 url 补 `https://` 前缀(本意是补全裸域名),而 JPEG 的 base64 载荷
|
||||
恒以 `/9j/` 开头——`urlparse("base64:///9j/...")` 的 netloc 为空,于是被补成
|
||||
`https://base64:///9j/...`,协议端拿 `base64` 当域名去 fetch,必报
|
||||
`ActionFailed ... fetch failed`(2026-10-10 实测:group_tools 的老婆图)。
|
||||
见 `repair_mangled_url`;它修的是**构造期就已损坏的值**,与发送方式无关,
|
||||
file 模式下同样执行。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -104,14 +116,44 @@ def _env_flag(name: str, default: bool) -> bool:
|
||||
return raw.strip().lower() not in {"", "0", "false", "no", "off"}
|
||||
|
||||
|
||||
#: 资源发送方式(显式配置,仿 HoshinoBot 的 RES_PROTOCOL)
|
||||
MEDIA_MODES: tuple[str, ...] = ("http", "base64", "file")
|
||||
|
||||
|
||||
def _mode() -> str:
|
||||
"""当前发送方式:`http`(默认)/ `base64` / `file`。
|
||||
|
||||
- `http` 改写为 `http://<可达地址>:<端口>/media/<token>`,协议端来拉;
|
||||
拿不到可达地址时回退 base64(媒体服务没挂上 / 无网卡)
|
||||
- `base64` 一律内联,字节由 bot 自己携带,不依赖协议端回拨
|
||||
- `file` 完全不改写,直接把本地路径交给适配器 —— 只有**协议端就在本机**
|
||||
时才是对的(跨机必 ENOENT,见模块文档)
|
||||
|
||||
配置键 `HEXI_MEDIA_MODE` 优先;没配时按旧的 `HEXI_INLINE_MEDIA` /
|
||||
`HEXI_INLINE_MEDIA_URL` 两个布尔开关推断,老 .env 行为不变。
|
||||
"""
|
||||
raw = (os.environ.get("HEXI_MEDIA_MODE") or "").strip().lower()
|
||||
if raw in MEDIA_MODES:
|
||||
return raw
|
||||
if raw:
|
||||
logger.warning(
|
||||
f"HEXI_MEDIA_MODE={raw!r} 不是 {'/'.join(MEDIA_MODES)} 之一,按默认 http 处理"
|
||||
)
|
||||
if not _env_flag("HEXI_INLINE_MEDIA", True):
|
||||
return "file"
|
||||
if not _env_flag("HEXI_INLINE_MEDIA_URL", True):
|
||||
return "base64"
|
||||
return "http"
|
||||
|
||||
|
||||
def _enabled() -> bool:
|
||||
"""总开关(默认开)。关闭的语义是「协议端就在本机」,此时 file:// 是有效的。"""
|
||||
return _env_flag("HEXI_INLINE_MEDIA", True)
|
||||
"""是否启用改写。`file` 模式=「协议端就在本机」,此时 file:// 本来就是对的。"""
|
||||
return _mode() != "file"
|
||||
|
||||
|
||||
def _url_enabled() -> bool:
|
||||
"""URL 优先模式(默认开);关掉则一律走 base64。"""
|
||||
return _env_flag("HEXI_INLINE_MEDIA_URL", True)
|
||||
"""是否用 URL(`/media/<token>`):只有 http 模式用,其余走 base64 / 原样。"""
|
||||
return _mode() == "http"
|
||||
|
||||
|
||||
def _ttl() -> float:
|
||||
@@ -187,12 +229,31 @@ def _is_private(ip: str) -> bool:
|
||||
return a == 10 or (a == 172 and 16 <= b <= 31) or (a == 192 and b == 168)
|
||||
|
||||
|
||||
def _route_ip() -> str | None:
|
||||
"""默认路由出口 IP —— 连外网时内核会挑的那张网卡。
|
||||
|
||||
多网卡机器只看网卡列表挑不出「协议端该往哪连」:VPN / 虚拟机 / WSL /
|
||||
EasyTier / Tailscale 的地址**同样是私有段**,排序时纯按字符串比大小,
|
||||
`10.0.0.129`(EasyTier 虚拟网卡)就会赢过真实局域网的 `192.168.2.186`,
|
||||
协议端拉不到 → `ActionFailed ... fetch failed`(2026-10-08 实测)。
|
||||
UDP connect 不发任何包,只是让内核回答"出口在哪";无网络时返回 None。
|
||||
"""
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
sock.connect(("8.8.8.8", 80))
|
||||
return str(sock.getsockname()[0])
|
||||
except OSError:
|
||||
return None
|
||||
finally:
|
||||
sock.close()
|
||||
|
||||
|
||||
def _local_ips() -> list[str]:
|
||||
"""本机可被协议端访问的 IPv4 列表(私有网段优先,排除环回/链路本地)。
|
||||
"""本机可被协议端访问的 IPv4 列表(默认路由网卡最前,其次私有网段优先)。
|
||||
|
||||
与 hexi/web_hub/dashboard.py 的 `_local_ips` 同源逻辑(那边供 /hub 首页
|
||||
「协议端配对」展示,即用户已经验证过能用的地址);此处刻意不 import,
|
||||
避免 hexi.core → hexi.web_hub 的反向依赖。
|
||||
避免 hexi.core → hexi.web_hub 的反向依赖。改这里记得同步那边。
|
||||
"""
|
||||
ips: list[str] = []
|
||||
try:
|
||||
@@ -223,6 +284,11 @@ def _local_ips() -> list[str]:
|
||||
if ip not in out:
|
||||
out.append(ip)
|
||||
out.sort(key=lambda ip: (not _is_private(ip), ip))
|
||||
route = _route_ip()
|
||||
if route in out:
|
||||
# 默认路由那张网卡排最前:它才是协议端大概率连得到的地址(见 _route_ip)
|
||||
out.remove(route)
|
||||
out.insert(0, route)
|
||||
return out
|
||||
|
||||
|
||||
@@ -288,7 +354,9 @@ def mount_media_endpoint() -> None:
|
||||
global _mounted
|
||||
|
||||
if not _enabled():
|
||||
logger.info("出站媒体内联化已禁用 (HEXI_INLINE_MEDIA=false),按原值发送")
|
||||
logger.info(
|
||||
f"出站媒体内联化已禁用 (HEXI_MEDIA_MODE={_mode()}),按原值发送"
|
||||
)
|
||||
return
|
||||
if _mounted:
|
||||
return
|
||||
@@ -349,6 +417,8 @@ async def _to_base64(path: Path) -> str:
|
||||
|
||||
async def rewrite_file(value: Any) -> str | None:
|
||||
"""改写单个 file 字段值;无需改写(或改写失败)返回 None。"""
|
||||
if not _enabled():
|
||||
return None # file 模式:路径原样交给适配器
|
||||
path = local_file_of(value)
|
||||
if path is None:
|
||||
return None
|
||||
@@ -386,6 +456,32 @@ async def rewrite_upload_file(value: Any) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
# alconna 的 Media.__post_init__ 会给解析不出 hostname 的 url 补 `https://`
|
||||
# 前缀(本意是补全 example.com/a.png 这类裸域名),而 OneBot file 字段认的
|
||||
# 两个伪协议都可能解析不出 hostname:
|
||||
# - `base64://` + JPEG 载荷(恒以 "/9j/" 开头)→ netloc 为空 → 变成
|
||||
# `https://base64:///9j/...`,协议端只会把 "base64" 当域名去 fetch,必死
|
||||
# - `file:///D:/x.jpg`(netloc 也为空)同理被补头
|
||||
# 带上 http(s) 外衣后这两个值不再有任何一方能解释,还原回伪协议本身。
|
||||
_MANGLED_SCHEMES = ("base64://", "file://")
|
||||
|
||||
|
||||
def repair_mangled_url(value: Any) -> str | None:
|
||||
"""还原被 alconna 补了 http(s):// 头的伪协议值;无需还原返回 None。
|
||||
|
||||
只认「http(s):// 前缀 + 伪协议」这种必然非法的组合——裸域名补成
|
||||
https:// 是 alconna 的正常行为,不受影响。
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
for prefix in ("https://", "http://"):
|
||||
if value.startswith(prefix):
|
||||
rest = value[len(prefix) :]
|
||||
if rest.startswith(_MANGLED_SCHEMES):
|
||||
return rest
|
||||
return None
|
||||
|
||||
|
||||
def _as_segment(node: Any) -> tuple[str, dict] | None:
|
||||
"""识别消息段对象 → (type, data);不是消息段返回 None。
|
||||
|
||||
@@ -404,7 +500,12 @@ def _as_segment(node: Any) -> tuple[str, dict] | None:
|
||||
|
||||
async def _walk_segment(seg_type: str, seg_data: dict, depth: int) -> None:
|
||||
if seg_type in MEDIA_SEG_TYPES:
|
||||
rewritten = await rewrite_file(seg_data.get("file"))
|
||||
value = seg_data.get("file")
|
||||
# 先还原构造期被 alconna 补坏的伪协议,再按需改写本地文件
|
||||
repaired = repair_mangled_url(value)
|
||||
if repaired is not None:
|
||||
value = seg_data["file"] = repaired
|
||||
rewritten = await rewrite_file(value)
|
||||
if rewritten is not None:
|
||||
seg_data["file"] = rewritten
|
||||
# 合并转发节点:content 可能是 str / list[dict] / Message,里面还可能嵌媒体
|
||||
@@ -455,9 +556,10 @@ async def _inline_outbound_media(bot: BaseBot, api: str, data: dict[str, Any]) -
|
||||
|
||||
钩子拿到的是 `_call_api` 之前**同一个可变 data dict**,且 hook 的 task group
|
||||
在 `_call_api` 之前 await 完成,所以原地改生效。
|
||||
|
||||
这里不按 `_enabled()` 提前退出:两个改写函数各自按发送方式把关(file 模式
|
||||
原样放行),而 `repair_mangled_url` 修的是构造期就坏掉的值,任何模式都该修。
|
||||
"""
|
||||
if not _enabled():
|
||||
return
|
||||
is_send = api in _SEND_APIS
|
||||
is_upload = api in _FILE_UPLOAD_APIS
|
||||
if not (is_send or is_upload) or not isinstance(data, dict):
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
"""定时任务的统一默认值与日志(`nonebot_plugin_apscheduler` 的薄封装)。
|
||||
|
||||
直接用 `scheduler.scheduled_job` / `scheduler.add_job` 时,各插件参数各写各的,
|
||||
而 APScheduler 的原生默认值有两个坑:
|
||||
|
||||
1. **不设时区** → 用服务器本地时区。换机器、进容器、系统时区不是
|
||||
Asia/Shanghai 时,`cron(hour=4)` 这类任务会整体错位 —— 而且错得很安静:
|
||||
任务照常在跑,只是时间不对(`interval` 类任务不受影响)。
|
||||
2. **`misfire_grace_time` 默认只有 1 秒** → 事件循环被卡一下(同步 IO、GC、
|
||||
笔记本休眠唤醒)任务就被判"错过"直接丢弃,一天少跑几次很难发现。
|
||||
|
||||
本模块只做两件事:**补默认值** + **包一层日志**(起止 DEBUG、异常 ERROR 后吞掉,
|
||||
不让坏任务刷爆 apscheduler 自己的日志)。调用方显式给的参数一律优先,
|
||||
触发条件一个字都不改。
|
||||
|
||||
用法(跟原来一模一样,只是换个导入):
|
||||
|
||||
from hexi.core.scheduler import add_job, scheduled_job
|
||||
|
||||
@scheduled_job("cron", hour=4, id="temp_cleanup")
|
||||
async def cleanup(): ...
|
||||
|
||||
add_job(fn, "interval", minutes=10, id="xxx", replace_existing=True)
|
||||
|
||||
动态删任务仍用 `nonebot_plugin_apscheduler.scheduler.remove_job(...)`。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from nonebot import require
|
||||
from nonebot.log import logger
|
||||
|
||||
require("nonebot_plugin_apscheduler")
|
||||
|
||||
from nonebot_plugin_apscheduler import scheduler # noqa: E402
|
||||
|
||||
#: 默认时区:显式钉死,不跟随服务器本地时区(见模块文档第 1 条)
|
||||
DEFAULT_TIMEZONE = "Asia/Shanghai"
|
||||
#: 默认宽限:宁可晚做也别不做(见模块文档第 2 条)
|
||||
DEFAULT_MISFIRE_GRACE_TIME = 60
|
||||
|
||||
#: 补进 kwargs 的默认值;`coalesce` 在 APScheduler 里本来就是 True,
|
||||
#: 写出来是为了表明"积压的任务只补跑一次"是有意为之
|
||||
_DEFAULTS: dict[str, Any] = {
|
||||
"timezone": DEFAULT_TIMEZONE,
|
||||
"misfire_grace_time": DEFAULT_MISFIRE_GRACE_TIME,
|
||||
"coalesce": True,
|
||||
}
|
||||
|
||||
|
||||
def _with_defaults(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
for key, value in _DEFAULTS.items():
|
||||
kwargs.setdefault(key, value)
|
||||
return kwargs
|
||||
|
||||
|
||||
async def _run_logged(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
||||
"""跑一个任务:起止 DEBUG、异常 ERROR 后吞掉(任务失败不该外抛给调度器)。"""
|
||||
logger.debug(f"定时任务开始: {func.__name__}")
|
||||
try:
|
||||
result = await func(*args, **kwargs)
|
||||
except Exception:
|
||||
logger.exception(f"定时任务异常: {func.__name__}")
|
||||
return None
|
||||
logger.debug(f"定时任务完成: {func.__name__}")
|
||||
return result
|
||||
|
||||
|
||||
def scheduled_job(*args: Any, **kwargs: Any) -> Callable[..., Any]:
|
||||
"""`scheduler.scheduled_job` + 统一默认值与日志(装饰器,用法不变)。"""
|
||||
_with_defaults(kwargs)
|
||||
|
||||
def deco(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
# functools.wraps 保住 __name__:APScheduler 不指定 id 时用函数名做任务 id
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*fargs: Any, **fkwargs: Any) -> Any:
|
||||
return await _run_logged(func, *fargs, **fkwargs)
|
||||
|
||||
return scheduler.scheduled_job(*args, **kwargs)(wrapper)
|
||||
|
||||
return deco
|
||||
|
||||
|
||||
def add_job(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
||||
"""`scheduler.add_job` + 统一默认值(动态增删任务用)。
|
||||
|
||||
注意包了一层:注册的是包装函数,`func.__name__` 由 `functools.wraps` 保住,
|
||||
所以不显式传 `id` 时的默认任务 id 仍是原函数名。删除任务请用原函数的
|
||||
id(或显式传 `id`)配合 `scheduler.remove_job(...)`。
|
||||
"""
|
||||
|
||||
@functools.wraps(func)
|
||||
async def wrapper(*fargs: Any, **fkwargs: Any) -> Any:
|
||||
return await _run_logged(func, *fargs, **fkwargs)
|
||||
|
||||
return scheduler.add_job(wrapper, *args, **_with_defaults(kwargs))
|
||||
@@ -0,0 +1,418 @@
|
||||
"""跨插件的「任务槽」:长任务的进度上报 + 同槽去重 + SSE 推送。
|
||||
|
||||
起因(2026-10-10):video_analysis 的 temp 清理是个几十秒级的阻塞任务,直接放在
|
||||
HTTP 请求里会把页面交互一起卡住,用户还可能连点几次。这类需求(清缓存、批量下载、
|
||||
批量生成…)会越来越多,所以在 core 出一份通用机制,插件按需注册。
|
||||
|
||||
三个概念:
|
||||
|
||||
- **槽(slot)**:字符串标识,约定 `插件id:动作`(如 `video_analysis:temp_cleanup`)。
|
||||
**同一槽同时只跑一个任务** —— 重复 `submit` 不新建,直接把你交给正在跑的那个并返回
|
||||
`reused=True`,前端据此禁用按钮 / 提示"已在运行"(服务端是最终防线,不能只靠前端)。
|
||||
- **执行**:任务体是**同步可调用对象**,由本模块丢线程池跑(`asyncio.to_thread`),
|
||||
所以清理/下载那类阻塞 IO 不会卡事件循环;函数签名固定为 `func(report) -> 结果`。
|
||||
- **进度**:`report(done=…, total=…, phase=…, message=…, warning=…)`,在**任务体所在的
|
||||
线程**里调用,内部用 `loop.call_soon_threadsafe` 切回事件循环改状态。
|
||||
推送按"百分比桶"节流(同一百分点内不重复发),上万个文件也不会灌爆 SSE。
|
||||
|
||||
`report` 的**返回值**是协作式取消信号:`False` 表示有人请求取消,任务体应尽快收尾
|
||||
(清理这种天然可中断的任务在循环里判断即可,不要硬杀线程)。
|
||||
|
||||
推送方式:`subscribe()` 拿一个 `asyncio.Queue`,此后**任何**任务变化都会往队列里塞一份
|
||||
全量快照;hub 的 `/api/tasks/stream` 直接搭在它上面。任务少、变化不频繁,全量快照比
|
||||
增量协议省心得多(代价是 `result` 之类的字段要克制,别往里塞大对象)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from collections.abc import Callable
|
||||
|
||||
from nonebot.log import logger
|
||||
|
||||
#: 任务体:接收 report 回调,返回任意结果(会放进 task.result)
|
||||
TaskFunc = Callable[[Callable[..., bool]], Any]
|
||||
|
||||
#: 终态任务保留时长(秒)与总量上限(超出先丢最旧的终态任务)
|
||||
TERMINAL_TTL = 30 * 60
|
||||
MAX_TASKS = 50
|
||||
|
||||
#: 状态取值:running → done / failed / canceled
|
||||
STATUS_RUNNING = "running"
|
||||
STATUS_DONE = "done"
|
||||
STATUS_FAILED = "failed"
|
||||
STATUS_CANCELED = "canceled"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Task:
|
||||
"""一条任务记录;`to_dict()` 就是 SSE / REST 的负载。
|
||||
|
||||
字段分四组:**身份**(id / slot / title / plugin_id)、**来源**(trigger / actor /
|
||||
group_id,出问题时能查到是谁点的)、**进度**(phase / done / total / unit /
|
||||
message,to_dict 里归成 `progress` 对象)、**结果**(result / warnings / error)
|
||||
与**时间**(created / started / updated / finished)。
|
||||
"""
|
||||
|
||||
# ── 身份 ──────────────────────────────────────────
|
||||
id: str
|
||||
#: 槽名:同槽互斥,约定 `插件id:动作`
|
||||
slot: str
|
||||
#: 展示名,如"清理 temp(保留 7 天)"
|
||||
title: str
|
||||
plugin_id: str = ""
|
||||
|
||||
# ── 来源(审计) ───────────────────────────────────
|
||||
#: 触发方式:web / command / schedule
|
||||
trigger: str = ""
|
||||
#: web 触发是登录用户名,command 触发是 QQ 号
|
||||
actor: str = ""
|
||||
#: command 触发时的群号
|
||||
group_id: str = ""
|
||||
|
||||
# ── 状态 ──────────────────────────────────────────
|
||||
status: str = STATUS_RUNNING
|
||||
#: 有人请求取消(任务体通过 `report()` 的返回值感知)
|
||||
cancel_requested: bool = False
|
||||
|
||||
# ── 进度 ──────────────────────────────────────────
|
||||
#: 阶段文案("扫描中" / "清理中" / "回收空目录")
|
||||
phase: str = ""
|
||||
done: int = 0
|
||||
#: 0 = 总数未知(扫描阶段没有进度条)
|
||||
total: int = 0
|
||||
#: 进度单位("个文件"),前端拼 "3/7 个文件"
|
||||
unit: str = ""
|
||||
#: 当前细节("已删除 2 个")
|
||||
message: str = ""
|
||||
|
||||
# ── 结果 ──────────────────────────────────────────
|
||||
result: Any = None
|
||||
#: 部分失败(跳过的被占用文件等)——不算失败,但要让人看见
|
||||
warnings: list[str] = field(default_factory=list)
|
||||
error: str = ""
|
||||
|
||||
# ── 时间 ──────────────────────────────────────────
|
||||
created_at: float = field(default_factory=time.time)
|
||||
started_at: float = 0.0
|
||||
updated_at: float = field(default_factory=time.time)
|
||||
finished_at: float = 0.0
|
||||
#: 结束信号(不出 API):命令层要"等任务跑完再回复"时用 `TASKS.wait(task)`
|
||||
_done: asyncio.Event = field(
|
||||
default_factory=asyncio.Event, repr=False, compare=False
|
||||
)
|
||||
|
||||
@property
|
||||
def active(self) -> bool:
|
||||
return self.status == STATUS_RUNNING
|
||||
|
||||
@property
|
||||
def percent(self) -> int | None:
|
||||
"""0-100;总数未知时 None(前端显示不确定进度)。"""
|
||||
if not self.active:
|
||||
return 100
|
||||
if self.total <= 0:
|
||||
return None
|
||||
return max(0, min(100, int(self.done * 100 / self.total)))
|
||||
|
||||
@property
|
||||
def duration(self) -> float:
|
||||
"""耗时(秒):跑着就是"已运行多久",结束就是总用时,没开始是 0。"""
|
||||
if not self.started_at:
|
||||
return 0.0
|
||||
end = self.finished_at or (time.time() if self.active else 0.0)
|
||||
return max(0.0, end - self.started_at)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"id": self.id,
|
||||
"slot": self.slot,
|
||||
"title": self.title,
|
||||
"plugin_id": self.plugin_id,
|
||||
"trigger": self.trigger,
|
||||
"actor": self.actor,
|
||||
"group_id": self.group_id,
|
||||
"status": self.status,
|
||||
"active": self.active,
|
||||
"cancel_requested": self.cancel_requested,
|
||||
"progress": {
|
||||
"percent": self.percent,
|
||||
"done": self.done,
|
||||
"total": self.total,
|
||||
"unit": self.unit,
|
||||
"phase": self.phase,
|
||||
"message": self.message,
|
||||
},
|
||||
"result": self.result,
|
||||
"warnings": list(self.warnings),
|
||||
"error": self.error,
|
||||
"created_at": self.created_at,
|
||||
"started_at": self.started_at,
|
||||
"updated_at": self.updated_at,
|
||||
"finished_at": self.finished_at,
|
||||
"duration": self.duration,
|
||||
}
|
||||
|
||||
|
||||
class TaskSlots:
|
||||
"""任务注册表(进程内单例 `TASKS`)。
|
||||
|
||||
只在事件循环里用:`submit` 需要 running loop 来建后台任务,`subscribe` 的队列
|
||||
也绑定在事件循环上。任务体本身跑在线程池里,但改状态一律切回事件循环。
|
||||
"""
|
||||
|
||||
def __init__(self, *, max_tasks: int = MAX_TASKS, ttl: float = TERMINAL_TTL) -> None:
|
||||
self._tasks: dict[str, Task] = {}
|
||||
self._subscribers: set[asyncio.Queue[list[dict[str, Any]]]] = set()
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._bg: set[asyncio.Task[None]] = set()
|
||||
self.max_tasks = max_tasks
|
||||
self.ttl = ttl
|
||||
#: 节流用:task.id -> 上次推送的百分比桶(-1 表示"未定进度")
|
||||
self._last_bucket: dict[str, int] = {}
|
||||
|
||||
# ── 查询 ──────────────────────────────────────────────
|
||||
|
||||
def list(self) -> list[Task]:
|
||||
"""全部任务,新的在前(顺带回收超龄/超量的终态任务)。"""
|
||||
self._prune()
|
||||
return sorted(self._tasks.values(), key=lambda t: t.created_at, reverse=True)
|
||||
|
||||
def get(self, task_id: str) -> Task | None:
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
def active(self, slot: str) -> Task | None:
|
||||
"""该槽正在跑的任务(没有则 None)—— 去重判定就是它。"""
|
||||
for task in self._tasks.values():
|
||||
if task.slot == slot and task.active:
|
||||
return task
|
||||
return None
|
||||
|
||||
# ── 提交 / 取消 ───────────────────────────────────────
|
||||
|
||||
def submit(
|
||||
self,
|
||||
*,
|
||||
slot: str,
|
||||
title: str,
|
||||
func: TaskFunc,
|
||||
plugin_id: str = "",
|
||||
trigger: str = "",
|
||||
actor: str = "",
|
||||
group_id: str = "",
|
||||
) -> tuple[Task, bool]:
|
||||
"""提交一个任务;同槽已有在跑的 → 不新建,返回 (已有任务, True)。
|
||||
|
||||
立即返回(不等任务跑完),任务体在线程池里执行。
|
||||
"""
|
||||
running = self.active(slot)
|
||||
if running is not None:
|
||||
logger.info(f"任务槽 {slot} 已有任务在跑({running.title}),本次提交复用")
|
||||
return running, True
|
||||
|
||||
self._loop = asyncio.get_running_loop()
|
||||
task = Task(
|
||||
id=uuid.uuid4().hex[:12],
|
||||
slot=slot,
|
||||
title=title,
|
||||
plugin_id=plugin_id,
|
||||
trigger=trigger,
|
||||
actor=str(actor or ""),
|
||||
group_id=str(group_id or ""),
|
||||
)
|
||||
self._tasks[task.id] = task
|
||||
self._prune()
|
||||
# 审计信息进日志:任务表本身是内存态(重启即清),但"谁在什么时候点了什么"
|
||||
# 要能在 /hub 日志页翻到
|
||||
who = trigger or "unknown"
|
||||
if actor:
|
||||
who += f" · {actor}"
|
||||
if group_id:
|
||||
who += f"@群{group_id}"
|
||||
logger.info(f"任务槽 {slot} 启动任务:{title}({who})")
|
||||
|
||||
bg = asyncio.create_task(self._run(task, func))
|
||||
self._bg.add(bg)
|
||||
bg.add_done_callback(self._bg.discard)
|
||||
|
||||
self._publish()
|
||||
return task, False
|
||||
|
||||
def cancel(self, task_id: str) -> bool:
|
||||
"""请求取消:只置标志,**任务体自己决定何时收尾**(协作式,不硬杀线程)。
|
||||
|
||||
任务体在 `report(...)` 的返回值里看到 False 就该尽快结束并返回已有结果,
|
||||
任务随即落到 `canceled` 状态(结果与 warnings 保留,方便看"取消前清了多少")。
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if task is None or not task.active:
|
||||
return False
|
||||
task.cancel_requested = True
|
||||
task.updated_at = time.time()
|
||||
logger.info(f"任务槽 {task.slot} 收到取消请求:{task.title}")
|
||||
self._publish()
|
||||
return True
|
||||
|
||||
async def wait(self, task: Task) -> Task:
|
||||
"""等任务跑完并返回它(已结束则立即返回)。
|
||||
|
||||
命令层用:"跑完再回复结果";Web 层不用(提交完就返回,进度走 SSE)。
|
||||
"""
|
||||
await task._done.wait()
|
||||
return task
|
||||
|
||||
def dismiss(self, task_id: str) -> bool:
|
||||
"""从列表里删掉一条**已结束**的任务,返回是否删掉了。
|
||||
|
||||
在跑的任务删不得(删了就再也没人回传它的状态了),只能先取消、等它收尾;
|
||||
所以这里对 active 的任务直接返回 False,让调用方去提示。
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if task is None or task.active:
|
||||
return False
|
||||
self._tasks.pop(task_id, None)
|
||||
self._last_bucket.pop(task_id, None)
|
||||
self._publish()
|
||||
return True
|
||||
|
||||
def clear_finished(self) -> int:
|
||||
"""清掉所有终态任务(在跑的不动),返回清掉的条数。"""
|
||||
finished = [tid for tid, t in self._tasks.items() if not t.active]
|
||||
for tid in finished:
|
||||
self._tasks.pop(tid, None)
|
||||
self._last_bucket.pop(tid, None)
|
||||
if finished:
|
||||
self._publish()
|
||||
return len(finished)
|
||||
|
||||
# ── 订阅(SSE) ───────────────────────────────────────
|
||||
|
||||
def subscribe(self) -> asyncio.Queue[list[dict[str, Any]]]:
|
||||
queue: asyncio.Queue[list[dict[str, Any]]] = asyncio.Queue()
|
||||
self._subscribers.add(queue)
|
||||
return queue
|
||||
|
||||
def unsubscribe(self, queue: asyncio.Queue[list[dict[str, Any]]]) -> None:
|
||||
self._subscribers.discard(queue)
|
||||
|
||||
def snapshot(self) -> list[dict[str, Any]]:
|
||||
return [t.to_dict() for t in self.list()]
|
||||
|
||||
# ── 内部:执行与推送 ──────────────────────────────────
|
||||
|
||||
async def _run(self, task: Task, func: TaskFunc) -> None:
|
||||
"""线程池里跑任务体;进度回调切回事件循环;异常落进 task.error。"""
|
||||
task.started_at = time.time()
|
||||
task.updated_at = task.started_at
|
||||
|
||||
def report(
|
||||
done: int | None = None,
|
||||
total: int | None = None,
|
||||
phase: str | None = None,
|
||||
unit: str | None = None,
|
||||
message: str | None = None,
|
||||
warning: str | None = None,
|
||||
) -> bool:
|
||||
"""任务体所在线程调用;返回 False 表示"请尽快收尾"(协作式取消)。"""
|
||||
loop = self._loop
|
||||
if loop is not None:
|
||||
loop.call_soon_threadsafe(
|
||||
self._apply_progress, task, done, total, phase, unit, message, warning
|
||||
)
|
||||
# cancel_requested 由事件循环线程写、工作线程读:CPython 下读 bool 是安全的
|
||||
return not task.cancel_requested
|
||||
|
||||
try:
|
||||
result = await asyncio.to_thread(func, report)
|
||||
except Exception as e: # noqa: BLE001 - 任务边界:失败即任务状态,不外抛
|
||||
logger.exception(f"任务失败({task.slot}):{task.title}")
|
||||
task.status = STATUS_FAILED
|
||||
task.error = str(e) or type(e).__name__
|
||||
else:
|
||||
task.status = STATUS_CANCELED if task.cancel_requested else STATUS_DONE
|
||||
task.result = result
|
||||
warn = f",{len(task.warnings)} 条警告" if task.warnings else ""
|
||||
logger.info(
|
||||
f"任务{'取消' if task.status == STATUS_CANCELED else '完成'}"
|
||||
f"({task.slot}):{task.title},用时 {task.duration:.1f}s{warn}"
|
||||
)
|
||||
task.finished_at = time.time()
|
||||
task.updated_at = task.finished_at
|
||||
task._done.set()
|
||||
self._publish()
|
||||
|
||||
def _apply_progress(
|
||||
self,
|
||||
task: Task,
|
||||
done: int | None,
|
||||
total: int | None,
|
||||
phase: str | None,
|
||||
unit: str | None,
|
||||
message: str | None,
|
||||
warning: str | None = None,
|
||||
) -> None:
|
||||
"""在事件循环里更新进度;按百分比桶节流推送。"""
|
||||
if not task.active: # 任务已结束(被 clear / 已失败)→ 丢弃迟到的进度
|
||||
return
|
||||
if done is not None:
|
||||
task.done = max(0, int(done))
|
||||
if total is not None:
|
||||
task.total = max(0, int(total))
|
||||
if phase is not None:
|
||||
task.phase = phase
|
||||
if unit is not None:
|
||||
task.unit = unit
|
||||
if message is not None:
|
||||
task.message = message
|
||||
if warning:
|
||||
task.warnings.append(str(warning))
|
||||
task.updated_at = time.time()
|
||||
|
||||
bucket = task.percent if task.percent is not None else -1
|
||||
noteworthy = phase is not None or message is not None or bool(warning)
|
||||
if bucket != self._last_bucket.get(task.id) or noteworthy:
|
||||
self._last_bucket[task.id] = bucket
|
||||
self._publish()
|
||||
|
||||
def _publish(self) -> None:
|
||||
"""给所有订阅者塞一份全量快照(队列满则丢掉最旧的,宁可丢进度也不阻塞)。"""
|
||||
if not self._subscribers:
|
||||
return
|
||||
snapshot = self.snapshot()
|
||||
for queue in list(self._subscribers):
|
||||
if queue.full():
|
||||
try:
|
||||
queue.get_nowait()
|
||||
except asyncio.QueueEmpty: # pragma: no cover - full 与 empty 之间被取走
|
||||
pass
|
||||
try:
|
||||
queue.put_nowait(snapshot)
|
||||
except asyncio.QueueFull: # pragma: no cover - 上面已腾位
|
||||
continue
|
||||
|
||||
def _prune(self) -> None:
|
||||
"""先清超龄终态任务,再按上限砍最旧的终态任务(在跑的不动)。"""
|
||||
now = time.time()
|
||||
for tid, task in list(self._tasks.items()):
|
||||
if not task.active and now - task.updated_at > self.ttl:
|
||||
self._tasks.pop(tid, None)
|
||||
self._last_bucket.pop(tid, None)
|
||||
|
||||
overflow = len(self._tasks) - self.max_tasks
|
||||
if overflow <= 0:
|
||||
return
|
||||
finished = sorted(
|
||||
(t for t in self._tasks.values() if not t.active),
|
||||
key=lambda t: t.created_at,
|
||||
)
|
||||
for task in finished[:overflow]:
|
||||
self._tasks.pop(task.id, None)
|
||||
self._last_bucket.pop(task.id, None)
|
||||
|
||||
|
||||
#: 全局单例(hub 路由与各插件共用)
|
||||
TASKS = TaskSlots()
|
||||
Reference in new issue
Block a user