mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-28 16:20:56 +08:00
molanp
39ed1ade14
* Update version constraint for nonebot-plugin-htmlrender * Refactor nonebot initialization in cli.py Simplified nonebot initialization by removing redundant parameters. * 更新 requirements.txt
413 lines
13 KiB
Python
413 lines
13 KiB
Python
"""zx CLI — 绪山真寻 Bot 命令行工具
|
||
|
||
用法:
|
||
zx run 启动 launcher
|
||
zx run-worker 启动 worker（由 launcher 调用）
|
||
zx version 显示版本信息
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import atexit
|
||
import importlib.metadata
|
||
import os
|
||
from pathlib import Path
|
||
import signal
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
GRACEFUL_SHUTDOWN_TIMEOUT = 15
|
||
WORKER_POLL_INTERVAL = 0.1
|
||
RESTART_POLL_INTERVAL = 0.5
|
||
WORKER_SOFT_EXIT_TIMEOUT = 15.0
|
||
WORKER_TERMINATE_TIMEOUT = 5.0
|
||
WORKER_KILL_TIMEOUT = 5.0
|
||
ENV_EXAMPLE_FILE = ".env.example"
|
||
ENV_DEV_FILE = ".env.dev"
|
||
|
||
|
||
def _env_assignment_key(line: str, *, include_commented: bool = False) -> str | None:
|
||
stripped = line.strip()
|
||
if include_commented and stripped.startswith("#"):
|
||
stripped = stripped[1:].lstrip()
|
||
if not stripped or stripped.startswith("#") or "=" not in stripped:
|
||
return None
|
||
key = stripped.split("=", 1)[0].strip()
|
||
return key if key.replace("_", "").isalnum() else None
|
||
|
||
|
||
def _env_key(line: str) -> str | None:
|
||
return _env_assignment_key(line)
|
||
|
||
|
||
def _env_block_key(block: list[str]) -> str | None:
|
||
for line in block:
|
||
if key := _env_key(line):
|
||
return key
|
||
return None
|
||
|
||
|
||
def _env_block_anchor_key(block: list[str]) -> str | None:
|
||
for line in block:
|
||
if key := _env_assignment_key(line, include_commented=True):
|
||
return key
|
||
return None
|
||
|
||
|
||
def _split_env_blocks(lines: list[str]) -> list[tuple[int, list[str]]]:
|
||
blocks: list[tuple[int, list[str]]] = []
|
||
current: list[str] = []
|
||
start_index = 0
|
||
for index, line in enumerate(lines):
|
||
if line.strip():
|
||
if not current:
|
||
start_index = index
|
||
current.append(line)
|
||
elif current:
|
||
blocks.append((start_index, current))
|
||
current = []
|
||
|
||
if current:
|
||
blocks.append((start_index, current))
|
||
return blocks
|
||
|
||
|
||
def _find_env_block_start(lines: list[str], key: str) -> int | None:
|
||
for start_index, block in _split_env_blocks(lines):
|
||
if _env_block_anchor_key(block) == key:
|
||
return start_index
|
||
return None
|
||
|
||
|
||
def _insert_env_block_before(
|
||
lines: list[str],
|
||
index: int,
|
||
block: list[str],
|
||
) -> list[str]:
|
||
insert_block = block.copy()
|
||
if index > 0 and lines[index - 1].strip():
|
||
insert_block.insert(0, "\n")
|
||
if index < len(lines) and insert_block and insert_block[-1].strip():
|
||
insert_block.append("\n")
|
||
return lines[:index] + insert_block + lines[index:]
|
||
|
||
|
||
def _sync_env_missing_items(project_root: Path) -> None:
|
||
"""Copy missing .env keys from .env.example without touching existing values."""
|
||
example_path = project_root / ENV_EXAMPLE_FILE
|
||
env_path = project_root / ENV_DEV_FILE
|
||
if not example_path.exists():
|
||
return
|
||
if not env_path.exists():
|
||
env_path.write_text(example_path.read_text(encoding="utf-8"), encoding="utf-8")
|
||
_launcher_log("已根据 .env.example 生成 .env.dev")
|
||
return
|
||
|
||
example_lines = example_path.read_text(encoding="utf-8").splitlines(keepends=True)
|
||
env_lines = env_path.read_text(encoding="utf-8").splitlines(keepends=True)
|
||
example_blocks = _split_env_blocks(example_lines)
|
||
existing_keys = {key for line in env_lines if (key := _env_key(line))}
|
||
missing_blocks: list[tuple[int, list[str]]] = []
|
||
|
||
for block_index, (_, block) in enumerate(example_blocks):
|
||
key = _env_block_key(block)
|
||
if key and key not in existing_keys:
|
||
missing_blocks.append((block_index, block))
|
||
|
||
if not missing_blocks:
|
||
return
|
||
|
||
updated_lines = env_lines
|
||
added_keys: list[str] = []
|
||
for block_index, block in missing_blocks:
|
||
key = _env_block_key(block)
|
||
if not key:
|
||
continue
|
||
anchor_index = len(updated_lines)
|
||
for _, next_block in example_blocks[block_index + 1 :]:
|
||
next_key = _env_block_anchor_key(next_block)
|
||
if not next_key:
|
||
continue
|
||
if (found := _find_env_block_start(updated_lines, next_key)) is not None:
|
||
anchor_index = found
|
||
break
|
||
updated_lines = _insert_env_block_before(updated_lines, anchor_index, block)
|
||
existing_keys.add(key)
|
||
added_keys.append(key)
|
||
|
||
env_path.write_text("".join(updated_lines), encoding="utf-8")
|
||
_launcher_log(f"已补齐 .env.dev 缺失配置: {', '.join(added_keys)}")
|
||
|
||
|
||
def _launcher_log(message: str) -> None:
|
||
sys.stderr.write(f"[zx launcher] {message}\n")
|
||
sys.stderr.flush()
|
||
|
||
|
||
def _print_version() -> None:
|
||
try:
|
||
ver = importlib.metadata.version("zhenxun-bot")
|
||
except importlib.metadata.PackageNotFoundError:
|
||
ver = "unknown"
|
||
sys.stdout.write(f"zhenxun-bot {ver}\n")
|
||
|
||
|
||
def _ensure_project_root() -> Path:
|
||
cwd = Path.cwd()
|
||
if not (cwd / "zhenxun").is_dir():
|
||
sys.stderr.write("错误: 当前目录不是 zhenxun_bot 项目目录。\n")
|
||
sys.stderr.write("请在项目根目录（包含 zhenxun/ 目录的位置）执行 zx run。\n")
|
||
sys.exit(1)
|
||
|
||
cwd_str = str(cwd)
|
||
if cwd_str not in sys.path:
|
||
sys.path.insert(0, cwd_str)
|
||
return cwd
|
||
|
||
|
||
def _run_worker() -> None:
|
||
"""启动 Bot worker（必须在项目目录下执行）"""
|
||
project_root = _ensure_project_root()
|
||
_sync_env_missing_items(project_root)
|
||
|
||
import contextlib
|
||
import platform
|
||
|
||
import nonebot
|
||
|
||
htmlrender_browser_channel = None
|
||
system = platform.system()
|
||
|
||
if system == "Windows":
|
||
import winreg
|
||
|
||
paths = {
|
||
"chrome": r"SOFTWARE\Clients\StartMenuInternet\Google Chrome\DefaultIcon",
|
||
"msedge": r"SOFTWARE\Clients\StartMenuInternet\Microsoft Edge\DefaultIcon",
|
||
}
|
||
for name, path in paths.items():
|
||
with contextlib.suppress(FileNotFoundError):
|
||
winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, path)
|
||
htmlrender_browser_channel = name
|
||
break
|
||
|
||
elif system == "Darwin":
|
||
mac_paths = {
|
||
"chrome": "/Applications/Google Chrome.app",
|
||
"msedge": "/Applications/Microsoft Edge.app",
|
||
}
|
||
for name, path in mac_paths.items():
|
||
if Path(path).exists():
|
||
htmlrender_browser_channel = name
|
||
break
|
||
|
||
if htmlrender_browser_channel:
|
||
nonebot.logger.info(
|
||
f"使用 {htmlrender_browser_channel} 作为 htmlrender 驱动启动..."
|
||
)
|
||
|
||
nonebot.init(htmlrender_browser_channel=htmlrender_browser_channel)
|
||
|
||
from nonebot.adapters.onebot.v11 import Adapter as OneBotV11Adapter
|
||
|
||
from zhenxun.configs.config import BotConfig
|
||
|
||
driver = nonebot.get_driver()
|
||
driver.register_adapter(OneBotV11Adapter)
|
||
enabled_adapters = ["OneBot V11"]
|
||
|
||
if BotConfig.qq_adapter_load:
|
||
try:
|
||
from nonebot.adapters.qq import Adapter as QQAdapter # type: ignore
|
||
except ImportError as e:
|
||
raise RuntimeError(
|
||
"QQ_ADAPTER_LOAD=True 但未安装 nonebot-adapter-qq，"
|
||
"请安装后再开启 QQ 官方适配器。"
|
||
) from e
|
||
driver.register_adapter(QQAdapter)
|
||
enabled_adapters.append("QQ")
|
||
|
||
nonebot.logger.info(f"已启用适配器: {', '.join(enabled_adapters)}")
|
||
|
||
nonebot.load_plugins("zhenxun/builtin_plugins")
|
||
nonebot.load_plugins("zhenxun/plugins")
|
||
|
||
for ext in BotConfig.ext_path:
|
||
ext = ext.strip()
|
||
if ext:
|
||
nonebot.logger.info(f"加载第三方插件目录: {ext}")
|
||
nonebot.load_plugins(ext)
|
||
|
||
nonebot.run(timeout_graceful_shutdown=GRACEFUL_SHUTDOWN_TIMEOUT)
|
||
|
||
|
||
def _build_worker_command() -> list[str]:
|
||
return [sys.executable, "-m", "zhenxun.cli", "run-worker"]
|
||
|
||
|
||
def _get_worker_creationflags() -> int:
|
||
if os.name == "nt":
|
||
return getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
|
||
return 0
|
||
|
||
|
||
def _wait_worker_exit(proc: subprocess.Popen, timeout_seconds: float) -> bool:
|
||
deadline = time.monotonic() + timeout_seconds
|
||
while time.monotonic() < deadline:
|
||
if proc.poll() is not None:
|
||
return True
|
||
time.sleep(WORKER_POLL_INTERVAL)
|
||
return proc.poll() is not None
|
||
|
||
|
||
def _terminate_worker(proc: subprocess.Popen) -> None:
|
||
if proc.poll() is not None:
|
||
return
|
||
_launcher_log(f"stopping worker pid={proc.pid}")
|
||
if os.name == "nt":
|
||
ctrl_break_event = getattr(signal, "CTRL_BREAK_EVENT", None)
|
||
if ctrl_break_event is not None:
|
||
try:
|
||
_launcher_log(f"sending CTRL_BREAK_EVENT to worker pid={proc.pid}")
|
||
proc.send_signal(ctrl_break_event)
|
||
except Exception as e:
|
||
_launcher_log(f"failed to send CTRL_BREAK_EVENT: {e!r}")
|
||
else:
|
||
if _wait_worker_exit(proc, WORKER_SOFT_EXIT_TIMEOUT):
|
||
_launcher_log(
|
||
f"worker pid={proc.pid} exited after CTRL_BREAK_EVENT "
|
||
f"with code {proc.returncode}"
|
||
)
|
||
return
|
||
_launcher_log(
|
||
f"worker pid={proc.pid} did not exit after "
|
||
f"{WORKER_SOFT_EXIT_TIMEOUT:.0f}s"
|
||
)
|
||
if _wait_worker_exit(proc, 1.0):
|
||
return
|
||
try:
|
||
_launcher_log(f"terminating worker pid={proc.pid}")
|
||
proc.terminate()
|
||
except Exception as e:
|
||
_launcher_log(f"failed to terminate worker: {e!r}")
|
||
else:
|
||
if _wait_worker_exit(proc, WORKER_TERMINATE_TIMEOUT):
|
||
_launcher_log(
|
||
f"worker pid={proc.pid} exited after terminate with code "
|
||
f"{proc.returncode}"
|
||
)
|
||
return
|
||
_launcher_log(f"worker pid={proc.pid} did not exit after terminate timeout")
|
||
_launcher_log(f"killing worker pid={proc.pid}")
|
||
proc.kill()
|
||
proc.wait(timeout=WORKER_KILL_TIMEOUT)
|
||
|
||
|
||
def _run_launcher() -> None:
|
||
cwd = _ensure_project_root()
|
||
from zhenxun.utils.restart_state import (
|
||
clear_launcher_restart_signal,
|
||
consume_launcher_restart_signal,
|
||
)
|
||
|
||
clear_launcher_restart_signal()
|
||
current_worker: subprocess.Popen | None = None
|
||
stop_requested = False
|
||
stop_signal: int | None = None
|
||
|
||
def _cleanup_current_worker() -> None:
|
||
if current_worker is not None:
|
||
_terminate_worker(current_worker)
|
||
|
||
atexit.register(_cleanup_current_worker)
|
||
|
||
def _handle_launcher_signal(signum, _frame) -> None:
|
||
nonlocal stop_requested, stop_signal
|
||
if stop_requested:
|
||
_launcher_log(f"received signal {signum} while stopping, exiting launcher")
|
||
raise SystemExit(128 + int(signum))
|
||
stop_requested = True
|
||
stop_signal = int(signum)
|
||
_launcher_log(f"received signal {signum}, scheduling worker shutdown")
|
||
|
||
handled_signals = [signal.SIGINT]
|
||
if hasattr(signal, "SIGTERM"):
|
||
handled_signals.append(signal.SIGTERM)
|
||
if hasattr(signal, "SIGBREAK"):
|
||
handled_signals.append(signal.SIGBREAK)
|
||
for sig in handled_signals:
|
||
try:
|
||
signal.signal(sig, _handle_launcher_signal)
|
||
except Exception:
|
||
pass
|
||
|
||
while True:
|
||
if stop_requested:
|
||
raise SystemExit(128 + int(stop_signal or signal.SIGINT))
|
||
worker_env = os.environ.copy()
|
||
worker_env["ZHENXUN_LAUNCHER_PID"] = str(os.getpid())
|
||
worker = subprocess.Popen(
|
||
_build_worker_command(),
|
||
cwd=str(cwd),
|
||
creationflags=_get_worker_creationflags(),
|
||
env=worker_env,
|
||
)
|
||
current_worker = worker
|
||
restart_requested = False
|
||
return_code: int | None = None
|
||
next_restart_check = 0.0
|
||
try:
|
||
while True:
|
||
return_code = worker.poll()
|
||
if return_code is not None:
|
||
break
|
||
if stop_requested:
|
||
clear_launcher_restart_signal()
|
||
_terminate_worker(worker)
|
||
raise SystemExit(128 + int(stop_signal or signal.SIGINT))
|
||
now = time.monotonic()
|
||
if now >= next_restart_check:
|
||
next_restart_check = now + RESTART_POLL_INTERVAL
|
||
if consume_launcher_restart_signal():
|
||
restart_requested = True
|
||
_launcher_log(
|
||
"detected restart request, stopping current worker"
|
||
)
|
||
_terminate_worker(worker)
|
||
return_code = worker.poll()
|
||
break
|
||
time.sleep(WORKER_POLL_INTERVAL)
|
||
except KeyboardInterrupt:
|
||
clear_launcher_restart_signal()
|
||
_terminate_worker(worker)
|
||
return
|
||
finally:
|
||
if current_worker is worker:
|
||
current_worker = None
|
||
|
||
if restart_requested or consume_launcher_restart_signal():
|
||
continue
|
||
raise SystemExit(return_code if return_code is not None else 1)
|
||
|
||
|
||
def main() -> None:
|
||
args = sys.argv[1:]
|
||
|
||
if not args or args[0] == "run":
|
||
_run_launcher()
|
||
elif args[0] == "run-worker":
|
||
_run_worker()
|
||
elif args[0] == "version":
|
||
_print_version()
|
||
elif args[0] in ("-h", "--help", "help"):
|
||
sys.stdout.write((__doc__ or "") + "\n")
|
||
else:
|
||
sys.stderr.write(f"未知命令: {args[0]}\n")
|
||
sys.stderr.write((__doc__ or "") + "\n")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|