K2cr2O1
c3b3541694
- 将函数名从 handle_admin_command 改为 admin_command_handler 以保持命名一致性 - 将变量名 comm_prefixes 改为 command_prefixes 以提高可读性 - 重命名 full_cmd 为 command_parts 和 cmd_name 为 command_name 以明确用途 - 简化 WebSocket 相关代码，移除未使用的导入 - 优化 main.py 中的初始化逻辑和变量命名
140 lines
4.3 KiB
Python
140 lines
4.3 KiB
Python
"""
|
||
NEO Bot 主程序入口
|
||
|
||
负责启动 WebSocket 连接，初始化插件系统，并提供热重载功能。
|
||
"""
|
||
import asyncio
|
||
import os
|
||
import sys
|
||
import time
|
||
from watchdog.observers import Observer
|
||
from watchdog.events import FileSystemEventHandler
|
||
|
||
# 初始化日志系统，必须在其他 core 模块导入之前执行
|
||
from core.utils.logger import logger
|
||
|
||
from core.managers.admin_manager import admin_manager
|
||
from core.ws import WS
|
||
from core.managers.plugin_manager import load_all_plugins
|
||
from core.managers.redis_manager import redis_manager
|
||
from core.utils.executor import run_in_thread_pool, initialize_executor
|
||
from core.config_loader import global_config as config
|
||
|
||
# 将项目根目录添加到 sys.path
|
||
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
sys.path.insert(0, ROOT_DIR)
|
||
|
||
|
||
|
||
|
||
class PluginReloadHandler(FileSystemEventHandler):
|
||
"""
|
||
文件变更处理器，用于热重载插件
|
||
|
||
继承自 watchdog.events.FileSystemEventHandler，
|
||
监听 base_plugins 目录下的文件变化，并触发插件重载。
|
||
"""
|
||
def __init__(self, loop: asyncio.AbstractEventLoop):
|
||
"""
|
||
初始化处理器
|
||
|
||
设置冷却时间，并保存主事件循环的引用。
|
||
"""
|
||
self.loop = loop
|
||
self.last_reload_time = 0
|
||
self.cooldown = 1.0 # 冷却时间，防止短时间内多次重载
|
||
|
||
def on_any_event(self, file_system_event):
|
||
"""
|
||
处理所有文件事件
|
||
|
||
:param file_system_event: watchdog 事件对象
|
||
"""
|
||
if file_system_event.is_directory:
|
||
return
|
||
|
||
# 只监控 py 文件
|
||
if not file_system_event.src_path.endswith(".py"):
|
||
return
|
||
|
||
# 过滤掉一些临时文件
|
||
if "__pycache__" in file_system_event.src_path:
|
||
return
|
||
|
||
# 简单的防抖动
|
||
current_time = time.time()
|
||
if current_time - self.last_reload_time < self.cooldown:
|
||
return
|
||
|
||
self.last_reload_time = current_time
|
||
|
||
logger.info(f"检测到文件变更: {file_system_event.src_path}")
|
||
logger.info("正在重载插件...")
|
||
|
||
try:
|
||
# 使用线程安全的方式在主事件循环中运行异步的插件加载函数
|
||
asyncio.run_coroutine_threadsafe(run_in_thread_pool(load_all_plugins), self.loop)
|
||
logger.success("插件重载完成")
|
||
except Exception as e:
|
||
logger.exception(f"重载失败: {e}")
|
||
|
||
|
||
@logger.catch
|
||
async def main():
|
||
"""
|
||
主函数
|
||
|
||
1. 启动文件监控（热重载）
|
||
2. 初始化 WebSocket 客户端
|
||
3. 建立连接并保持运行
|
||
"""
|
||
# 首次加载插件
|
||
await run_in_thread_pool(load_all_plugins)
|
||
|
||
# 初始化 Redis 连接
|
||
await redis_manager.initialize()
|
||
|
||
# 初始化管理员管理器
|
||
await admin_manager.initialize()
|
||
|
||
# 启动文件监控
|
||
# 监控 plugins 目录
|
||
plugin_path = os.path.join(os.path.dirname(__file__), "plugins")
|
||
|
||
# 获取当前事件循环并传递给处理器
|
||
loop = asyncio.get_running_loop()
|
||
event_handler = PluginReloadHandler(loop)
|
||
observer = Observer()
|
||
|
||
if os.path.exists(plugin_path):
|
||
observer.schedule(event_handler, plugin_path, recursive=True)
|
||
observer.start()
|
||
logger.info(f"已启动插件热重载监控: {plugin_path}")
|
||
else:
|
||
logger.warning(f"插件目录不存在 {plugin_path}")
|
||
|
||
try:
|
||
websocket_client = WS()
|
||
|
||
# 初始化代码执行器
|
||
code_executor = initialize_executor(websocket_client, config)
|
||
websocket_client.bot.code_executor = code_executor # 将执行器实例附加到 bot.bot 对象上
|
||
|
||
# 启动代码执行器的后台 worker
|
||
logger.debug("[Main] 检查是否需要启动代码执行 Worker...")
|
||
if code_executor and code_executor.docker_client:
|
||
logger.info("[Main] Docker 连接成功，正在启动代码执行 Worker...")
|
||
asyncio.create_task(code_executor.worker())
|
||
else:
|
||
logger.warning("[Main] 未启动代码执行 Worker，因为 Docker 客户端未初始化或连接失败。")
|
||
|
||
await websocket_client.connect()
|
||
finally:
|
||
if observer.is_alive():
|
||
observer.stop()
|
||
observer.join()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|