diff --git a/hexi/core/__init__.py b/hexi/core/__init__.py index 159778a..08279e7 100644 --- a/hexi/core/__init__.py +++ b/hexi/core/__init__.py @@ -9,6 +9,7 @@ from . import ( # noqa: E402,F401 plugin_control, plugin_manager, rate_limit, + task_slots, ) diff --git a/hexi/core/outbound_media.py b/hexi/core/outbound_media.py index 842c0d1..611a028 100644 --- a/hexi/core/outbound_media.py +++ b/hexi/core/outbound_media.py @@ -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/`**,由协议端主动来拉。 地址取自 `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/`,协议端来拉; + 拿不到可达地址时回退 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/`):只有 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): diff --git a/hexi/core/scheduler.py b/hexi/core/scheduler.py new file mode 100644 index 0000000..4366f85 --- /dev/null +++ b/hexi/core/scheduler.py @@ -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)) diff --git a/hexi/core/task_slots.py b/hexi/core/task_slots.py new file mode 100644 index 0000000..2bec9d9 --- /dev/null +++ b/hexi/core/task_slots.py @@ -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() diff --git a/hexi/plugins/nonebot_plugin_learning_chat/handlers/learn.py b/hexi/plugins/nonebot_plugin_learning_chat/handlers/learn.py index 507e307..7c361fd 100644 --- a/hexi/plugins/nonebot_plugin_learning_chat/handlers/learn.py +++ b/hexi/plugins/nonebot_plugin_learning_chat/handlers/learn.py @@ -20,7 +20,7 @@ from nonebot.typing import T_State require("nonebot_plugin_orm") require("nonebot_plugin_apscheduler") -from nonebot_plugin_apscheduler import scheduler +from hexi.core.scheduler import scheduled_job from nonebot_plugin_orm import get_session @@ -77,7 +77,7 @@ async def _(event: GroupMessageEvent, answers=Arg("answers")): ) -@scheduler.scheduled_job("interval", minutes=3, misfire_grace_time=5) +@scheduled_job("interval", minutes=3, misfire_grace_time=5) async def speak_up(): if not config_manager.config.total_enable: return diff --git a/tests/test_outbound_media.py b/tests/test_outbound_media.py index dd48ff2..e7f1c42 100644 --- a/tests/test_outbound_media.py +++ b/tests/test_outbound_media.py @@ -113,6 +113,134 @@ def test_non_string_values_are_left_alone(om, value): assert om.local_file_of(value) is None +# ------------------------------------------------------------ 可达地址选择 + + +def _fake_psutil(monkeypatch, ips_by_iface: dict[str, list[str]]) -> None: + """把 psutil.net_if_addrs 换成假网卡表(`_local_ips` 在函数内 import psutil)。""" + import socket as _socket + import types + + fake = types.ModuleType("psutil") + fake.net_if_addrs = lambda: { + name: [ + types.SimpleNamespace(family=_socket.AF_INET, address=ip) for ip in ips + ] + for name, ips in ips_by_iface.items() + } + monkeypatch.setitem(sys.modules, "psutil", fake) + + +def test_route_ip_wins_over_virtual_adapter(om, monkeypatch): + """回归 2026-10-08:EasyTier 的 10.0.0.129 按字符串排序赢过真实局域网, + 协议端拉不到媒体(ActionFailed fetch failed)。默认路由网卡必须排最前。""" + _fake_psutil( + monkeypatch, + { + "EasyTier": ["10.0.0.129"], + "以太网": ["192.168.2.186"], + "WLAN": ["169.254.65.98"], # 链路本地,应被过滤 + }, + ) + monkeypatch.setattr(om, "_route_ip", lambda: "192.168.2.186") + + ips = om._local_ips() + assert ips[0] == "192.168.2.186" + assert ips == ["192.168.2.186", "10.0.0.129"] + + +def test_local_ips_without_route_falls_back_to_private_first(om, monkeypatch): + """探不到默认路由(无网络/异常)时退回原排序:私有段优先、其余按字符串。""" + _fake_psutil( + monkeypatch, + {"Ethernet": ["192.168.2.186"], "VPN": ["10.0.0.129"], "eth1": ["203.0.113.9"]}, + ) + monkeypatch.setattr(om, "_route_ip", lambda: None) + assert om._local_ips() == ["10.0.0.129", "192.168.2.186", "203.0.113.9"] + + +def test_local_ips_filters_loopback_and_ipv6(om, monkeypatch): + _fake_psutil( + monkeypatch, + {"Loopback": ["127.0.0.1", "::1"], "Ethernet": ["192.168.2.186", "fe80::1"]}, + ) + monkeypatch.setattr(om, "_route_ip", lambda: None) + assert om._local_ips() == ["192.168.2.186"] + + +def test_route_ip_absent_from_candidates_is_ignored(om, monkeypatch): + """路由出口不在网卡表里(异常环境)时不该崩,也不该把结果搞乱。""" + _fake_psutil(monkeypatch, {"Ethernet": ["192.168.2.186"]}) + monkeypatch.setattr(om, "_route_ip", lambda: "10.9.9.9") + assert om._local_ips() == ["192.168.2.186"] + + +def test_route_ip_smoke(om): + """真跑一次 UDP connect(不发包):有默认路由时应给出本机 IPv4。""" + ip = om._route_ip() + assert ip is None or (ip.count(".") == 3 and not ip.startswith("127.")) + + +# ------------------------------------------------------------ 发送方式 + + +def _clear_mode_env(monkeypatch) -> None: + for key in ("HEXI_MEDIA_MODE", "HEXI_INLINE_MEDIA", "HEXI_INLINE_MEDIA_URL"): + monkeypatch.delenv(key, raising=False) + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + # 默认 http(协议端跨机来拉) + ({}, "http"), + ({"HEXI_MEDIA_MODE": "http"}, "http"), + ({"HEXI_MEDIA_MODE": " BASE64 "}, "base64"), # 大小写/空白容忍 + ({"HEXI_MEDIA_MODE": "file"}, "file"), + # 新键没配时按旧开关推断(老 .env 行为不变) + ({"HEXI_INLINE_MEDIA": "false"}, "file"), + ({"HEXI_INLINE_MEDIA_URL": "0"}, "base64"), + ({"HEXI_INLINE_MEDIA": "false", "HEXI_INLINE_MEDIA_URL": "false"}, "file"), + # 新键优先于旧开关 + ({"HEXI_MEDIA_MODE": "http", "HEXI_INLINE_MEDIA": "false"}, "http"), + # 非法值退回旧开关/默认 + ({"HEXI_MEDIA_MODE": "ftp"}, "http"), + ({"HEXI_MEDIA_MODE": "ftp", "HEXI_INLINE_MEDIA_URL": "false"}, "base64"), + ], +) +def test_mode_resolution(om, monkeypatch, env, expected): + _clear_mode_env(monkeypatch) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert om._mode() == expected + + +async def test_file_mode_passes_through(om, media_file: Path, monkeypatch): + """file 模式(协议端就在本机)完全不改写:消息与文件上传都原样交给适配器。""" + monkeypatch.setenv("HEXI_MEDIA_MODE", "file") + monkeypatch.setattr(om, "_mounted", True) + monkeypatch.setenv("HEXI_MEDIA_BASE_URL", "http://192.168.2.15:39697") + assert await om.rewrite_file(str(media_file)) is None + assert await om.rewrite_upload_file(str(media_file)) is None + + +async def test_base64_mode_skips_url(om, media_file: Path, monkeypatch): + """base64 模式即使媒体服务可用也不发 URL(协议端拨不通本机的场合)。""" + monkeypatch.setenv("HEXI_MEDIA_MODE", "base64") + monkeypatch.setattr(om, "_mounted", True) + monkeypatch.setenv("HEXI_MEDIA_BASE_URL", "http://192.168.2.15:39697") + assert await om.rewrite_file(str(media_file)) == "base64://aGVsbG8=" + # 文件上传不做 base64(太大)→ base64 模式下只能原样放行 + assert await om.rewrite_upload_file(str(media_file)) is None + + +async def test_http_mode_still_falls_back_to_base64(om, media_file: Path, monkeypatch): + """http 模式拿不到可达地址时仍回退 base64(没挂上媒体服务/无网卡)。""" + monkeypatch.setenv("HEXI_MEDIA_MODE", "http") + monkeypatch.setattr(om, "_mounted", False) + assert await om.rewrite_file(str(media_file)) == "base64://aGVsbG8=" + + # ------------------------------------------------------------ 改写目标 @@ -278,6 +406,67 @@ async def test_walk_respects_depth_limit(om, media_file: Path, monkeypatch): assert deepest["data"]["file"] == str(media_file) # 没被改写 +# ------------------------------------------------------------ 伪协议还原 + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + # JPEG 的 base64 载荷恒以 /9j/ 开头,alconna 构造 Image(url=...) 时 + # 会因 netloc 为空补上 https:// 前缀(2026-10-10 实测 fetch failed 的根因) + ("https://base64:///9j/4AAQSkZJRg", "base64:///9j/4AAQSkZJRg"), + ("http://base64:///9j/4AAQSkZJRg", "base64:///9j/4AAQSkZJRg"), + ("https://file:///D:/x.jpg", "file:///D:/x.jpg"), + ], +) +def test_repair_mangled_pseudo_scheme(om, value: str, expected: str): + assert om.repair_mangled_url(value) == expected + + +@pytest.mark.parametrize( + "value", + [ + # 裸域名补 https:// 是 alconna 的本意,正常链接一律不碰 + "https://example.com/a.jpg", + # 主机名恰好是 base64 的正常链接不能误伤(只有 base64:// 才算伪协议) + "https://base64.com/a.png", + # 没被补头的原值无需还原 + "base64:///9j/4AAQSkZJRg", + "file:///D:/x.jpg", + None, + 123, + ], +) +def test_repair_leaves_normal_values_alone(om, value): + assert om.repair_mangled_url(value) is None + + +async def test_walk_repairs_mangled_base64(om, monkeypatch): + """补头后的 base64 值进了 file 字段,walk 必须还原回伪协议本身。""" + monkeypatch.setattr(om, "_mounted", False) + seg = {"type": "image", "data": {"file": "https://base64:///9j/4AAQSkZJRg"}} + await om._walk(seg) + assert seg["data"]["file"] == "base64:///9j/4AAQSkZJRg" + + +async def test_walk_repairs_real_jpeg_base64(om, tmp_path: Path, monkeypatch): + """用真实 JPEG 魔数验证前提:载荷以 /9j/ 开头 → urlparse 解析不出 hostname + → 正是 alconna 会补 https:// 的情形;还原后应与原始 base64:// 串一致。""" + import base64 as _b64 + from urllib.parse import urlparse + + fp = tmp_path / "a.jpg" + fp.write_bytes(b"\xff\xd8\xff\xe0" + b"\x00" * 20) + payload = _b64.b64encode(fp.read_bytes()).decode() + assert payload.startswith("/9j/") # JPEG 恒如此,补头的前提 + assert urlparse(f"base64://{payload}").hostname is None # netloc 为空才会被补头 + + monkeypatch.setattr(om, "_mounted", False) + seg = {"type": "image", "data": {"file": f"https://base64://{payload}"}} + await om._walk(seg) + assert seg["data"]["file"] == f"base64://{payload}" + + # ------------------------------------------------------------ 钩子 @@ -316,6 +505,29 @@ async def test_hook_disabled_by_env(om, media_file: Path, monkeypatch): assert before.startswith("file://") +async def test_hook_repairs_mangled_value_even_in_file_mode( + om, media_file: Path, monkeypatch +): + """file 模式不改写本地文件,但构造期就坏掉的伪协议任何模式都该修。 + + 2026-10-10 老婆图报 `ActionFailed ... fetch failed` 发的就是这种值: + 与协议端在哪台机器无关,不还原则必死。 + """ + monkeypatch.setenv("HEXI_MEDIA_MODE", "file") + data = { + "message": Message( + [ + MessageSegment.image(media_file), + MessageSegment.image("https://base64:///9j/4AAQSkZJRg"), + ] + ) + } + before = data["message"][0].data["file"] + await om._inline_outbound_media(_FakeBot(), "send_group_msg", data) + assert data["message"][0].data["file"] == before # 本地文件原样 + assert data["message"][1].data["file"] == "base64:///9j/4AAQSkZJRg" + + async def test_hook_handles_forward_messages_key(om, media_file: Path, monkeypatch): monkeypatch.setattr(om, "_mounted", False) data = { diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..35ed019 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,161 @@ +"""定时任务薄封装(hexi/core/scheduler.py)单元测试 + +只测"补默认值 + 日志包装"这层:不启动 APScheduler、不初始化 NoneBot。 +做法是往 sys.modules 塞一个假的 `nonebot_plugin_apscheduler`(记录调用参数), +再把 `nonebot.require` 换成 no-op(模块 import 时会 require 那个插件), +然后按文件路径加载目标模块(与本仓其它 hexi/core 测试同法)。 + +默认值本身是这次改动的核心: +- `timezone` 不设 → 跟随服务器本地时区,`cron(hour=4)` 会静默错位 +- `misfire_grace_time` APScheduler 原生只有 1 秒 → 事件循环卡一下任务就被丢 +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +_MODULE_PATH = ( + Path(__file__).resolve().parents[1] / "hexi" / "core" / "scheduler.py" +) +_MOD_NAME = "hexi_core_scheduler_under_test" + + +class _FakeScheduler: + """记录注册参数;装饰器语义与 APScheduler 一致(返回被装饰的函数)。""" + + def __init__(self) -> None: + self.calls: list[tuple[str, tuple, dict, object]] = [] + + def scheduled_job(self, *args, **kwargs): + def deco(func): + self.calls.append(("scheduled_job", args, kwargs, func)) + return func + + return deco + + def add_job(self, func, *args, **kwargs): + self.calls.append(("add_job", args, kwargs, func)) + return func + + +class _FakeLogger: + """替代 loguru logger:记录 debug/exception 调用。""" + + def __init__(self) -> None: + self.debugs: list[str] = [] + self.exceptions: list[str] = [] + + def debug(self, msg: str, *a, **kw) -> None: + self.debugs.append(msg) + + def exception(self, msg: str, *a, **kw) -> None: + self.exceptions.append(msg) + + +@pytest.fixture +def sched(monkeypatch): + """加载 scheduler.py:假 apscheduler + no-op require。""" + fake_scheduler = _FakeScheduler() + fake_plugin = types.ModuleType("nonebot_plugin_apscheduler") + fake_plugin.scheduler = fake_scheduler + monkeypatch.setitem(sys.modules, "nonebot_plugin_apscheduler", fake_plugin) + + import nonebot + + monkeypatch.setattr(nonebot, "require", lambda name: None, raising=False) + + spec = importlib.util.spec_from_file_location(_MOD_NAME, _MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[_MOD_NAME] = module + spec.loader.exec_module(module) + + module.logger = _FakeLogger() # 日志也换成可观察的 + module._fake = fake_scheduler + + yield module + sys.modules.pop(_MOD_NAME, None) + + +# ───────────────────────── 默认值 ───────────────────────── + + +def test_scheduled_job_fills_defaults(sched): + @sched.scheduled_job("cron", hour=4, id="cleanup") + async def cleanup(): # pragma: no cover - 不执行 + pass + + _, args, kwargs, _ = sched._fake.calls[-1] + assert args == ("cron",) + assert kwargs["hour"] == 4 + assert kwargs["id"] == "cleanup" + assert kwargs["timezone"] == sched.DEFAULT_TIMEZONE == "Asia/Shanghai" + assert kwargs["misfire_grace_time"] == sched.DEFAULT_MISFIRE_GRACE_TIME == 60 + assert kwargs["coalesce"] is True + + +def test_scheduled_job_keeps_explicit_values(sched): + """调用方显式给的一律优先(如 learning_chat 的 5 秒宽限)。""" + + @sched.scheduled_job("interval", minutes=3, misfire_grace_time=5, timezone="UTC") + async def job(): # pragma: no cover + pass + + _, _, kwargs, _ = sched._fake.calls[-1] + assert kwargs["misfire_grace_time"] == 5 + assert kwargs["timezone"] == "UTC" + assert kwargs["minutes"] == 3 + + +def test_add_job_fills_defaults(sched): + """动态注册(video_analysis 的每日 cron 走这条)同样补默认值。""" + + async def cleanup(): # pragma: no cover + pass + + sched.add_job( + cleanup, + "cron", + hour=4, + id="video_analysis_temp_cleanup", + misfire_grace_time=3600, + replace_existing=True, + ) + + kind, args, kwargs, registered = sched._fake.calls[-1] + assert kind == "add_job" + assert args == ("cron",) + assert kwargs["timezone"] == "Asia/Shanghai" # 显式没给 → 补上(这次改动的重点) + assert kwargs["misfire_grace_time"] == 3600 # 显式给了 → 不覆盖 + assert kwargs["replace_existing"] is True + # 包了一层,但函数名保住(不传 id 时 APScheduler 用 __name__ 作任务 id) + assert registered is not cleanup + assert registered.__name__ == "cleanup" + + +# ───────────────────────── 包装行为 ───────────────────────── + + +async def test_wrapper_runs_and_returns_value(sched): + async def job(x): # noqa: ANN001 + return x * 2 + + registered = sched.scheduled_job("interval", minutes=1)(job) + assert registered.__name__ == "job" + assert await registered(21) == 42 + assert any("定时任务开始: job" in m for m in sched.logger.debugs) + assert any("定时任务完成: job" in m for m in sched.logger.debugs) + + +async def test_wrapper_swallows_exception(sched): + """任务抛异常 → 记日志并吞掉(不让坏任务外抛给调度器)。""" + + async def boom(): + raise RuntimeError("炸了") + + registered = sched.scheduled_job("interval", minutes=1)(boom) + assert await registered() is None + assert sched.logger.exceptions and "定时任务异常: boom" in sched.logger.exceptions[0] + assert not any("完成" in m for m in sched.logger.debugs) diff --git a/tests/test_task_slots.py b/tests/test_task_slots.py new file mode 100644 index 0000000..38433e9 --- /dev/null +++ b/tests/test_task_slots.py @@ -0,0 +1,347 @@ +"""任务槽(hexi/core/task_slots.py)单元测试 + +只测注册表本身:执行/去重/进度/取消/推送/回收。模块只依赖 asyncio + nonebot.log, +按路径裸加载即可(与 test_rate_limit 同法)。 + +任务体跑在线程池里,进度回调是"切回事件循环再改状态",所以断言进度时必须等 +状态落地 —— 用 `_until()` 轮询事件循环,别用 `sleep` 猜时序。 +""" + +import asyncio +import importlib.util +import sys +import threading +import time +from pathlib import Path + +import pytest + +_MODULE_PATH = Path(__file__).resolve().parents[1] / "hexi" / "core" / "task_slots.py" +_MOD_NAME = "hexi_core_task_slots_under_test" + + +@pytest.fixture(scope="module") +def ts(): + spec = importlib.util.spec_from_file_location(_MOD_NAME, _MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[_MOD_NAME] = module + spec.loader.exec_module(module) + yield module + sys.modules.pop(_MOD_NAME, None) + + +@pytest.fixture +def reg(ts): + """每条用例一个干净的注册表(不用全局单例,避免用例互相污染)。""" + return ts.TaskSlots() + + +async def _until(pred, timeout: float = 2.0) -> None: + """轮询事件循环直到条件成立(进度是跨线程/跨回调落地的,不能靠猜时序)。""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return + await asyncio.sleep(0.01) + raise AssertionError("等待条件超时") + + +# ───────────────────────── 执行与结果 ───────────────────────── + + +async def test_runs_in_thread_and_records_result(reg): + def work(report): + return "done!" + + task, reused = reg.submit(slot="plug:job", title="干活", func=work, plugin_id="plug") + assert reused is False + assert task.active # submit 即返回,此时后台任务还没被调度 + + await asyncio.wait_for(reg.wait(task), 2) + assert task.status == "done" + assert task.result == "done!" + assert task.started_at > 0 + assert task.finished_at >= task.started_at + assert task.duration >= 0 + assert reg.active("plug:job") is None + + +async def test_exception_marks_failed_without_raising(reg): + def boom(report): + raise RuntimeError("炸了") + + task, _ = reg.submit(slot="plug:boom", title="会炸", func=boom) + await asyncio.wait_for(reg.wait(task), 2) + assert task.status == "failed" + assert "炸了" in task.error + assert task.percent == 100 # 终态一律显示满(前端别再转圈) + + +# ───────────────────────── 同槽去重 ───────────────────────── + + +async def test_same_slot_reuses_running_task(reg): + calls = [] + gate = threading.Event() + + def slow(report): + calls.append(1) + gate.wait(2) + return "ok" + + first, reused = reg.submit(slot="plug:clean", title="清理", func=slow) + second, reused2 = reg.submit(slot="plug:clean", title="清理(再点一次)", func=slow) + + assert reused is False and reused2 is True + assert second.id == first.id # 返回的是正在跑的那个 + gate.set() + await asyncio.wait_for(reg.wait(first), 2) + assert calls == [1] # 任务体只跑了一次 + + +async def test_same_slot_available_again_after_finish(reg): + def work(report): + return 1 + + first, _ = reg.submit(slot="plug:clean", title="第一次", func=work) + await asyncio.wait_for(reg.wait(first), 2) + + second, reused = reg.submit(slot="plug:clean", title="第二次", func=work) + assert reused is False and second.id != first.id + await asyncio.wait_for(reg.wait(second), 2) + + +async def test_different_slots_run_in_parallel(reg): + gate = threading.Event() + + def slow(report): + gate.wait(2) + + a, _ = reg.submit(slot="plug:a", title="A", func=slow) + b, reused = reg.submit(slot="plug:b", title="B", func=slow) + assert reused is False and b.id != a.id + gate.set() + await asyncio.gather(reg.wait(a), reg.wait(b)) + + +# ───────────────────────── 进度 ───────────────────────── + + +async def test_progress_fields_and_percent(reg): + gate = threading.Event() + + def work(report): + report(done=3, total=12, phase="清理中", unit="个文件", message="已删除 2 个") + gate.wait(2) + return "ok" + + task, _ = reg.submit(slot="plug:job", title="进度", func=work) + await _until(lambda: task.done == 3) + + assert (task.total, task.unit, task.phase) == (12, "个文件", "清理中") + assert task.message == "已删除 2 个" + assert task.percent == 25 + + payload = task.to_dict()["progress"] + assert payload == { + "percent": 25, "done": 3, "total": 12, "unit": "个文件", + "phase": "清理中", "message": "已删除 2 个", + } + + gate.set() + await asyncio.wait_for(reg.wait(task), 2) + + +async def test_unknown_total_reports_none_percent(reg): + """扫描阶段总数未知 → percent None(前端显示不确定进度而不是 0%)。""" + gate = threading.Event() + + def work(report): + report(phase="扫描中", unit="个文件", message="正在统计文件…") + gate.wait(2) + + task, _ = reg.submit(slot="plug:job", title="扫描", func=work) + await _until(lambda: task.phase == "扫描中") + assert task.percent is None + gate.set() + await asyncio.wait_for(reg.wait(task), 2) + + +async def test_warnings_accumulate(reg): + def work(report): + report(warning="跳过 a.mp4(被占用)") + report(warning="跳过 b.mp4(被占用)") + return (0, 2) + + task, _ = reg.submit(slot="plug:job", title="部分失败", func=work) + await asyncio.wait_for(reg.wait(task), 2) + assert task.to_dict()["warnings"] == ["跳过 a.mp4(被占用)", "跳过 b.mp4(被占用)"] + + +async def test_progress_is_throttled(reg): + """上千次上报不能变成上千次推送(按百分比桶 + 只在有意义的字段变化时发)。""" + queue = reg.subscribe() + count = 500 + + def work(report): + for i in range(1, count + 1): + report(done=i, total=count) + return "ok" + + task, _ = reg.submit(slot="plug:job", title="节流", func=work) + await asyncio.wait_for(reg.wait(task), 2) + + frames = 0 + while not queue.empty(): + queue.get_nowait() + frames += 1 + # 理论上限:每个百分点一帧 + 起止几帧;远小于 500 次上报 + assert frames <= 110, f"推送 {frames} 帧,节流没生效" + assert frames >= 2 # 但也不能一帧都不发 + reg.unsubscribe(queue) + + +# ───────────────────────── 协作式取消 ───────────────────────── + + +async def test_cancel_stops_cooperative_task(reg): + def work(report): + for i in range(1, 1000): + if not report(done=i, total=999): + return ("canceled", i) # 任务体自己收尾,带上已有进度 + time.sleep(0.005) + return ("finished", 999) + + task, _ = reg.submit(slot="plug:job", title="可取消", func=work) + await _until(lambda: task.done >= 1) + + assert reg.cancel(task.id) is True + assert task.cancel_requested is True + + await asyncio.wait_for(reg.wait(task), 2) + assert task.status == "canceled" + assert task.result[0] == "canceled" + # 取消不是硬杀:任务体返回前的进度照常保留 + assert task.done < 999 + + +async def test_cancel_finished_or_unknown_task_is_noop(reg): + def work(report): + return 1 + + task, _ = reg.submit(slot="plug:job", title="快活", func=work) + await asyncio.wait_for(reg.wait(task), 2) + assert reg.cancel(task.id) is False + assert reg.cancel("不存在的 id") is False + + +# ───────────────────────── 订阅 / 回收 ───────────────────────── + + +async def test_subscribe_receives_snapshots(reg): + queue = reg.subscribe() + assert queue.empty() + + def work(report): + report(done=1, total=2, phase="干活") + return "ok" + + task, _ = reg.submit(slot="plug:job", title="推送", func=work) + await asyncio.wait_for(reg.wait(task), 2) + + assert not queue.empty() + frames = [] + while not queue.empty(): + frames.append(queue.get_nowait()) + assert frames[0][0]["title"] == "推送" + assert frames[-1][0]["status"] == "done" + # 每帧都是全量快照 + assert isinstance(frames[-1], list) + + reg.unsubscribe(queue) + reg.unsubscribe(queue) # 幂等 + + +async def test_clear_finished_keeps_running(reg): + gate = threading.Event() + + def slow(report): + gate.wait(2) + + def quick(report): + return 1 + + running, _ = reg.submit(slot="plug:slow", title="在跑", func=slow) + done, _ = reg.submit(slot="plug:quick", title="跑完了", func=quick) + await asyncio.wait_for(reg.wait(done), 2) + + assert reg.clear_finished() == 1 + assert reg.get(done.id) is None + assert reg.get(running.id) is not None + + gate.set() + await asyncio.wait_for(reg.wait(running), 2) + + +async def test_prune_by_max_tasks_keeps_running(reg): + small = type(reg)() # 同类型新实例 + small.max_tasks = 2 + + gate = threading.Event() + + def slow(report): + gate.wait(2) + + def quick(report): + return 1 + + running, _ = small.submit(slot="s:run", title="在跑", func=slow) + for i in range(3): + t, _ = small.submit(slot=f"s:{i}", title=f"第{i}次", func=quick) + await asyncio.wait_for(small.wait(t), 2) + + assert len(small.list()) <= 2 # 上限生效(终态的被砍) + assert small.get(running.id) is not None # 在跑的绝不动 + gate.set() + await asyncio.wait_for(small.wait(running), 2) + + +async def test_dismiss_only_finished(reg): + """已结束的能删,在跑的删不掉(删了状态就没人回传了)。""" + gate = threading.Event() + + def slow(report): + gate.wait(2) + + def quick(report): + return 1 + + running, _ = reg.submit(slot="plug:slow", title="在跑", func=slow) + done, _ = reg.submit(slot="plug:quick", title="跑完了", func=quick) + await asyncio.wait_for(reg.wait(done), 2) + + assert reg.dismiss(running.id) is False # 在跑:拒绝 + assert reg.get(running.id) is not None + assert reg.dismiss(done.id) is True + assert reg.get(done.id) is None + assert reg.dismiss(done.id) is False # 已删过 → 再删是 no-op + assert reg.dismiss("不存在的 id") is False + + gate.set() + await asyncio.wait_for(reg.wait(running), 2) + + +async def test_audit_fields_round_trip(reg): + def work(report): + return "ok" + + task, _ = reg.submit( + slot="plug:job", title="审计", func=work, + plugin_id="video_analysis", trigger="web", actor="admin", group_id="123456", + ) + await asyncio.wait_for(reg.wait(task), 2) + d = task.to_dict() + assert (d["plugin_id"], d["trigger"], d["actor"], d["group_id"]) == ( + "video_analysis", "web", "admin", "123456", + ) + assert d["status"] == "done" and d["active"] is False