zhenxun_bot/zhenxun/builtin_plugins/hooks/call_hook.py

import asyncio
from collections import deque
from typing import Any

from nonebot.adapters import Bot, Message
from nonebot.adapters.onebot.v11 import MessageSegment
from nonebot_plugin_apscheduler import scheduler

from zhenxun.configs.config import Config
from zhenxun.models.bot_message_store import BotMessageStore
from zhenxun.services.log import logger
from zhenxun.utils.enum import BotSentType
from zhenxun.utils.manager.message_manager import MessageManager
from zhenxun.utils.platform import PlatformUtils

LOG_COMMAND = "MessageHook"


_BOT_MSG_BUFFER: deque[dict[str, Any]] = deque()
_BOT_MSG_BUFFER_LOCK = asyncio.Lock()
_BOT_MSG_BULK_SIZE = 50
_PENDING_TASKS: set[asyncio.Task] = set()


async def _flush_bot_messages():
    async with _BOT_MSG_BUFFER_LOCK:
        if not _BOT_MSG_BUFFER:
            return
        items: list[dict[str, Any]] = []
        while _BOT_MSG_BUFFER:
            items.append(_BOT_MSG_BUFFER.popleft())
    try:
        await BotMessageStore.bulk_create([BotMessageStore(**it) for it in items])
    except Exception as e:
        logger.warning("批量写入BotMessageStore失败", LOG_COMMAND, e=e)
        # 尝试降级逐条写入，避免数据全部丢失
        try:
            for it in items:
                await BotMessageStore.create(**it)
        except Exception as e2:
            logger.warning("逐条写入BotMessageStore失败", LOG_COMMAND, e=e2)


async def _enqueue_bot_message(item: dict[str, Any]):
    async with _BOT_MSG_BUFFER_LOCK:
        _BOT_MSG_BUFFER.append(item)
        if len(_BOT_MSG_BUFFER) >= _BOT_MSG_BULK_SIZE:
            task = asyncio.create_task(_flush_bot_messages())
            _PENDING_TASKS.add(task)
            task.add_done_callback(_PENDING_TASKS.discard)


@scheduler.scheduled_job("interval", seconds=10)
async def _flush_bot_messages_job():
    await _flush_bot_messages()


def replace_message(message: Message) -> str:
    """将消息中的at、image、record、face替换为字符串

    参数:
        message: Message

    返回:
        str: 文本消息
    """
    result = ""
    for msg in message:
        if isinstance(msg, str):
            result += msg
        elif msg.type == "at":
            result += f"@{msg.data['qq']}"
        elif msg.type == "image":
            result += "[image]"
        elif msg.type == "record":
            result += "[record]"
        elif msg.type == "face":
            result += f"[face:{msg.data['id']}]"
        elif msg.type == "reply":
            result += ""
        else:
            result += str(msg)
    return result


def format_message_for_log(message: Message) -> str:
    """
    将消息对象转换为适合日志记录的字符串，对base64等长内容进行摘要处理。
    """
    if not isinstance(message, Message):
        return str(message)

    log_parts = []
    for seg in message:
        seg: MessageSegment
        if seg.type == "text":
            log_parts.append(seg.data.get("text", ""))
        elif seg.type in ("image", "record", "video"):
            file_info = seg.data.get("file", "")
            if isinstance(file_info, str) and file_info.startswith("base64://"):
                b64_data = file_info[9:]
                data_size_bytes = (len(b64_data) * 3) / 4 - b64_data.count("=", -2)
                log_parts.append(
                    f"[{seg.type}: base64, size={data_size_bytes / 1024:.2f}KB]"
                )
            else:
                log_parts.append(f"[{seg.type}]")
        elif seg.type == "at":
            log_parts.append(f"[@{seg.data.get('qq', 'unknown')}]")
        else:
            log_parts.append(f"[{seg.type}]")
    return "".join(log_parts)


@Bot.on_called_api
async def handle_api_result(
    bot: Bot, exception: Exception | None, api: str, data: dict[str, Any], result: Any
):
    if exception or api != "send_msg":
        return
    user_id = data.get("user_id")
    group_id = data.get("group_id")
    message_id = result.get("message_id")
    message: Message = data.get("message", "")
    message_type = data.get("message_type")
    try:
        # 记录消息id
        if user_id and message_id:
            MessageManager.add(str(user_id), str(message_id))
            logger.debug(
                f"收集消息id，user_id: {user_id}, msg_id: {message_id}", LOG_COMMAND
            )
    except Exception as e:
        logger.warning(
            f"收集消息id发生错误...data: {data}, result: {result}", LOG_COMMAND, e=e
        )
    if not Config.get_config("hook", "RECORD_BOT_SENT_MESSAGES"):
        return
    try:
        await _enqueue_bot_message(
            {
                "bot_id": bot.self_id,
                "user_id": user_id,
                "group_id": group_id,
                "sent_type": BotSentType.GROUP
                if message_type == "group"
                else BotSentType.PRIVATE,
                "text": replace_message(message),
                "plain_text": message.extract_plain_text()
                if isinstance(message, Message)
                else replace_message(message),
                "platform": PlatformUtils.get_platform(bot),
            }
        )
        logger.debug(f"消息发送记录，message: {format_message_for_log(message)}")
    except Exception as e:
        logger.warning(
            f"消息发送记录发生错误...data: {data}, result: {result}",
            LOG_COMMAND,
            e=e,
        )