bot.py

from core import check_python_version  # skipcq

check_python_version()  # noqa

from core.constants import config_path, config_filename  # skipcq

if not (config_path / config_filename).exists():
    import core.scripts.config_generate  # noqa

import asyncio
import importlib
import multiprocessing
import os
import shutil
import sys
import time
import traceback
from pathlib import Path

from dotenv import load_dotenv
from loguru import logger
from tortoise import Tortoise, run_async

from core.constants import ascii_art, bots_path, logs_path  # skipcq
from core.database import close_db
from core.config import CONFIG_READONLY_ENV


load_dotenv()
os.environ.setdefault("PYTHONIOENCODING", "UTF-8")
os.environ.setdefault("PYTHONPATH", str(Path(".").resolve()))


# Basic logger setup
try:
    logger.remove(0)
except ValueError:
    pass

Logger = logger.bind(name="BotDaemon")

logger_format = (
    "<cyan>[BotDaemon]</cyan>"
    "<yellow>[{name}:{function}:{line}]</yellow>"
    "<green>[{time:YYYY-MM-DD HH:mm:ss}]</green>"
    "<level>[{level}]:{message}</level>"
)
Logger.add(
    sys.stdout, format=logger_format, colorize=True, filter=lambda record: record["extra"].get("name") == "BotDaemon"
)

Logger.add(
    sink=logs_path / "BotDaemon_debug_{time:YYYY-MM-DD}.log",
    format=logger_format,
    rotation="00:00",
    retention="1 day",
    level="DEBUG",
    filter=lambda record: record["level"].name == "DEBUG" and record["extra"].get("name") == "BotDaemon",
    encoding="utf8",
)
Logger.add(
    sink=logs_path / "BotDaemon_{time:YYYY-MM-DD}.log",
    format=logger_format,
    rotation="00:00",
    retention="10 days",
    level="INFO",
    encoding="utf8",
    filter=lambda record: record["extra"].get("name") == "BotDaemon",
)


class RestartBot(Exception):
    pass


failed_to_start_attempts = {}
disabled_bots = []
processes: list[multiprocessing.Process] = []


def pre_init():
    Logger.info(ascii_art)
    Logger.info("Akaribot is launching...")
    from core.constants.path import cache_path

    if cache_path.exists():
        shutil.rmtree(cache_path)
    cache_path.mkdir(parents=True, exist_ok=True)

    Logger.info("Generating config...")

    # CoreConfig 的导入须留在函数内：multiprocessing 以 spawn / forkserver 启动子进程时
    # 会以 __mp_main__ 重新导入主模块，置于顶层将使每个子进程再次触发配置模板的生成。
    from core.config.base import CoreConfig
    from core.config.scan import scan_config_templates
    from core.constants.version import database_version
    from core.database.link import get_db_link
    from core.database.models import SenderUnionInfo, DBVersion

    # 配置的生成集中在此完成：子进程一律只读，此处遗漏的键将在子进程读取时抛出异常，
    # 因此任何模板加载失败均须中止启动，而非延至运行期才暴露。
    failed_templates = scan_config_templates()
    if failed_templates:
        Logger.critical(f"Failed to load config templates: {failed_templates}. Aborting.")
        sys.exit(1)

    if CoreConfig.debug:
        Logger.debug("Debug mode is enabled.")

    async def update_db():
        Logger.info("Checking database...")
        await Tortoise.init(db_url=get_db_link(), modules={"models": ["core.database.models"]})
        await Tortoise.generate_schemas(safe=True)

        Logger.info("Verifying database version...")

        query_dbver = await DBVersion.all().first()
        if not query_dbver:
            Logger.warning("Database version not found, converting old database...")
            from core.scripts.convert_database import convert_database

            await close_db()
            await convert_database()
            Logger.success("Database converted successfully!")
        elif query_dbver.version < database_version:
            Logger.info(f"Updating database from {query_dbver.version} to {database_version}...")
            from core.database.update import update_database

            await close_db()
            await update_database()
            Logger.success("Database updated successfully!")
        else:
            await close_db()

        Logger.success("Database check completed successfully!")

        Logger.info("Granting base superuser permission...")
        base_superuser = CoreConfig.base_superuser
        if base_superuser:
            if isinstance(base_superuser, str):
                base_superuser = [base_superuser]
            await Tortoise.init(db_url=get_db_link(), modules={"models": ["core.database.models"]})
            for bu in base_superuser:
                sender_info = await SenderUnionInfo.resolve_union(bu)
                await sender_info.edit_attr("superuser", True)
            await close_db()
            Logger.success("Base superuser permission granted successfully!")
        else:
            Logger.warning("The base superuser is not found, please setup it in the config file.")

    run_async(update_db())
    Logger.success("Pre-init completed successfully!")


def multiprocess_run_until_complete(func):
    mp = multiprocessing.get_context("spawn" if sys.platform in ["win32", "darwin"] else "forkserver")
    # pre_init 是唯一被允许写入配置的进程，须清除 RestartBot 重启循环中残留的只读标记
    os.environ.pop(CONFIG_READONLY_ENV, None)
    p = mp.Process(target=func, daemon=True)
    p.start()

    while True:
        if not p.is_alive():
            break
        time.sleep(1)
    # Process.close() 之后再访问 exitcode 会抛 ValueError，故须在 terminate_process 之前取值
    exitcode = p.exitcode
    terminate_process(p)
    if exitcode != 0:
        Logger.critical(f"Pre-init failed with exit code {exitcode}, aborting.")
        sys.exit(exitcode)


def go(bot_name: str, subprocess: bool = False, binary_mode: bool = False):
    from core.constants import Info

    Logger.info(f"[{bot_name}] Here we go!")
    Info.subprocess = subprocess
    Info.binary_mode = binary_mode

    try:
        importlib.import_module(f"bots.{bot_name}.bot")
    except Exception:
        # 适配器通常在导入阶段完成 SDK 登录；不能只捕获 ModuleNotFoundError，
        # 否则认证、网络和 SDK 初始化异常只会出现在子进程 stderr 中，守护进程日志看不到原因。
        Logger.error(f"[{bot_name}] Failed to start.\n{traceback.format_exc()}")
        sys.exit(1)


async def cleanup_tasks():
    loop = asyncio.get_event_loop()
    asyncio.all_tasks(loop=loop)


binary_mode = not sys.argv[0].endswith(".py")


async def run_bot():
    from core.config import CFGManager
    from core.server.run import run_async as server_run_async

    # 自此起 spawn 出的子进程一律只读：配置的生成已在 pre_init 中完成。
    # 须在任何 mp.Process 之前置位，spawn 会继承环境；restart_bot_process() 后续重启子进程时同样适用。
    os.environ[CONFIG_READONLY_ENV] = "1"

    mp = multiprocessing.get_context("spawn" if sys.platform in ["win32", "darwin"] else "forkserver")

    def restart_bot_process(bot_name: str):
        if (
            bot_name not in failed_to_start_attempts
            or time.time() - failed_to_start_attempts[bot_name]["timestamp"] > 60
        ):
            failed_to_start_attempts[bot_name] = {}
            failed_to_start_attempts[bot_name]["count"] = 0
            failed_to_start_attempts[bot_name]["timestamp"] = time.time()
        failed_to_start_attempts[bot_name]["count"] += 1
        failed_to_start_attempts[bot_name]["timestamp"] = time.time()
        if failed_to_start_attempts[bot_name]["count"] >= 3:
            Logger.error(f"Bot {bot_name} failed to start 3 times, abort to restart, please check the log.")
            return

        Logger.warning(f"Restarting bot {bot_name}...")
        p = mp.Process(
            target=go,
            args=(
                bot_name,
                True,
                binary_mode,
            ),
            name=bot_name,
            daemon=True,
        )
        p.start()
        processes.append(p)

    bots_list = [p.name for p in bots_path.iterdir() if p.is_dir() and not p.name.startswith("_")]

    for t in CFGManager.values:
        if t.startswith("bot_") and not t.endswith("_secret") and t[4:] in bots_list:
            if "enable" in CFGManager.values[t][t]:
                if not CFGManager.values[t][t]["enable"]:
                    disabled_bots.append(t[4:])
                    Logger.warning(f"Bot {t[4:]} is disabled in config, skip to launch.")
            else:
                Logger.warning(f'Bot {t[4:]} cannot found config "enable".')
                disabled_bots.append(t[4:])

    for bl in bots_list:
        if bl in disabled_bots:
            continue
        p = mp.Process(target=go, args=(bl, True, binary_mode), name=bl, daemon=True)
        p.start()
        processes.append(p)

    # run the server process
    server_process = mp.Process(target=server_run_async, args=(True, binary_mode), name="server", daemon=True)

    server_process.start()
    processes.append(server_process)

    while len(processes) > 1:
        for p in processes:
            if p.is_alive():
                continue
            if p.name == "server":
                if p.exitcode == 0:
                    sys.exit(0)
                if p.exitcode == 233:
                    Logger.warning(f"Process {p.pid} (server) exited with code 233, restart all bots.")
                    raise RestartBot
                Logger.critical(f"Process {p.pid} (server) exited with code {p.exitcode}, please check the log.")
                sys.exit(p.exitcode)
            if p.exitcode == 0:
                Logger.warning(f"Process {p.pid} ({p.name}) exited with code 0, abort to restart.")
                processes.remove(p)
                terminate_process(p)
                break
            if p.exitcode == 233:
                Logger.warning(f"Process {p.pid} ({p.name}) exited with code 233, restart all bots.")
                raise RestartBot
            Logger.critical(f"Process {p.pid} ({p.name}) exited with code {p.exitcode}, please check the log.")
            processes.remove(p)
            terminate_process(p)
            restart_bot_process(p.name)
            break
        await asyncio.sleep(1)
    Logger.critical("All bots exited unexpectedly, please check the output.")
    sys.exit(1)


def terminate_process(process: multiprocessing.Process):
    process.kill()
    process.join()
    process.close()


async def main_async():
    try:
        multiprocess_run_until_complete(pre_init)
        await run_bot()  # Process will block here so
    except RestartBot as e:
        for ps in processes:
            Logger.warning(f"Terminating process {ps.pid} ({ps.name})...")
            terminate_process(ps)
        processes.clear()
        raise e
    except (KeyboardInterrupt, SystemExit) as e:
        for ps in processes:
            terminate_process(ps)
        processes.clear()
        raise e
    except Exception as e:
        Logger.critical("An error occurred, please check the output.")
        traceback.print_exc()
        raise e
    finally:
        await close_db()


def main():
    # Capture the base import lists to avoid clearing essential modules when restarting
    def clear_import_cache():
        for m in list(sys.modules):
            if m not in base_import_lists:
                del sys.modules[m]

    base_import_lists = list(sys.modules)

    while True:
        try:
            asyncio.run(main_async())
        except RestartBot:
            clear_import_cache()
            continue
        except (KeyboardInterrupt, SystemExit):
            break


if __name__ == "__main__":
    # Detect if the program is already running
    lock_file_path = Path("./.bot.lock").resolve()
    if sys.platform == "win32":
        import msvcrt

        lock_file = open(lock_file_path, "w")  # skipcq
        try:
            msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
        except OSError:
            Logger.critical("Another instance is already running. Aborting.")
            sys.exit(1)
    else:
        import fcntl

        lock_file = open(lock_file_path, "w")  # skipcq
        try:
            fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            Logger.critical("Another instance is already running. Aborting.")
            sys.exit(1)

    main()