✨ 多机器人支持 (#179)

* 维护一个群号和QQ号与机器人实例的表

* 随机获取一个对应的机器人

* 获取所有机器人所在群的并集

* 监听到相关事件的时候刷新整个缓存

* 在发送错误后自动刷新缓存

* 合并相同类型的函数

* 添加测试

src/plugins/nonebot_bison/admin_page/api.py

@@ -4,7 +4,6 @@ from fastapi.exceptions import HTTPException
 from fastapi.param_functions import Depends
 from fastapi.routing import APIRouter
 from fastapi.security.oauth2 import OAuth2PasswordBearer
-from nonebot.adapters.onebot.v11.bot import Bot
 
 from ..apis import check_sub_target
 from ..config import (
@@ -16,7 +15,8 @@ from ..config import (
 from ..config.db_config import SubscribeDupException
 from ..platform import platform_manager
 from ..types import Target as T_Target
-from ..types import WeightConfig
+from ..types import User, WeightConfig
+from ..utils.get_bot import get_bot, get_groups
 from .jwt import load_jwt, pack_jwt
 from .token_manager import token_manager
 from .types import (
@@ -72,12 +72,13 @@ async def get_global_conf() -> GlobalConf:
 
 
 async def get_admin_groups(qq: int):
-    bot = nonebot.get_bot()
-    groups = await bot.call_api("get_group_list")
     res = []
-    for group in groups:
+    for group in await get_groups():
         group_id = group["group_id"]
-        users = await bot.call_api("get_group_member_list", group_id=group_id)
+        bot = get_bot(User(group_id, "group"))
+        if not bot:
+            continue
+        users = await bot.get_group_member_list(group_id=group_id)
         for user in users:
             if user["user_id"] == qq and user["role"] in ("owner", "admin"):
                 res.append({"id": group_id, "name": group["group_name"]})
@@ -88,9 +89,6 @@ async def get_admin_groups(qq: int):
 async def auth(token: str) -> TokenResp:
     if qq_tuple := token_manager.get_user(token):
         qq, nickname = qq_tuple
-        bot = nonebot.get_bot()
-        assert isinstance(bot, Bot)
-        groups = await bot.call_api("get_group_list")
         if str(qq) in nonebot.get_driver().config.superusers:
             jwt_obj = {
                 "id": qq,
@@ -101,7 +99,7 @@ async def auth(token: str) -> TokenResp:
                             "id": info["group_id"],
                             "name": info["group_name"],
                         },
-                        groups,
+                        await get_groups(),
                     )
                 ),
             }

src/plugins/nonebot_bison/scheduler/scheduler.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
 from typing import Optional, Type
 
-import nonebot
 from nonebot.adapters.onebot.v11.bot import Bot
 from nonebot.log import logger
 
@@ -10,6 +9,7 @@ from ..platform import platform_manager
 from ..send import send_msgs
 from ..types import Target
 from ..utils import ProcessContext, SchedulerConfig
+from ..utils.get_bot import get_bot
 from .aps import aps
 
 
@@ -105,9 +105,9 @@ class Scheduler:
 
         if not to_send:
             return
-        bot = nonebot.get_bot()
-        assert isinstance(bot, Bot)
+
         for user, send_list in to_send:
+            bot = get_bot(user)
             for send_post in send_list:
                 logger.info("send to {}: {}".format(user, send_post))
                 if not bot:

src/plugins/nonebot_bison/send.py

@@ -2,10 +2,12 @@ import time
 from typing import Literal, Union
 
 from nonebot.adapters.onebot.v11.bot import Bot
-from nonebot.adapters.onebot.v11.message import Message, MessageSegment
+from nonebot.adapters.onebot.v11.exception import ActionFailed
+from nonebot.adapters.onebot.v11.message import Message
 from nonebot.log import logger
 
 from .plugin_config import plugin_config
+from .utils.get_bot import refresh_bots
 
 QUEUE: list[
     tuple[
@@ -25,12 +27,16 @@ async def _do_send(
     user_type: Literal["group", "private", "group-forward"],
     msg: Union[str, Message],
 ):
-    if user_type == "group":
-        await bot.send_group_msg(group_id=user, message=msg)
-    elif user_type == "private":
-        await bot.send_private_msg(user_id=user, message=msg)
-    elif user_type == "group-forward":
-        await bot.send_group_forward_msg(group_id=user, messages=msg)
+    try:
+        if user_type == "group":
+            await bot.send_group_msg(group_id=user, message=msg)
+        elif user_type == "private":
+            await bot.send_private_msg(user_id=user, message=msg)
+        elif user_type == "group-forward":
+            await bot.send_group_forward_msg(group_id=user, messages=msg)
+    except ActionFailed:
+        await refresh_bots()
+        logger.warning(f"send msg failed, refresh bots")
 
 
 async def do_send_msgs():

src/plugins/nonebot_bison/utils/get_bot.py

@@ -0,0 +1,106 @@
+""" 提供获取 Bot 的方法 """
+import random
+from typing import Any, Optional
+
+import nonebot
+from nonebot import get_driver, on_notice
+from nonebot.adapters.onebot.v11 import (
+    Bot,
+    FriendAddNoticeEvent,
+    GroupDecreaseNoticeEvent,
+    GroupIncreaseNoticeEvent,
+)
+
+from ..types import User
+
+GROUP: dict[int, list[Bot]] = {}
+USER: dict[int, list[Bot]] = {}
+
+
+def get_bots() -> list[Bot]:
+    """获取所有 OneBot 11 Bot"""
+    bots = []
+    for bot in nonebot.get_bots().values():
+        if isinstance(bot, Bot):
+            bots.append(bot)
+    return bots
+
+
+async def refresh_bots():
+    """刷新缓存的 Bot 数据"""
+    GROUP.clear()
+    USER.clear()
+    for bot in get_bots():
+        # 获取群列表
+        groups = await bot.get_group_list()
+        for group in groups:
+            group_id = group["group_id"]
+            if group_id not in GROUP:
+                GROUP[group_id] = [bot]
+            else:
+                GROUP[group_id].append(bot)
+
+        # 获取好友列表
+        users = await bot.get_friend_list()
+        for user in users:
+            user_id = user["user_id"]
+            if user_id not in USER:
+                USER[user_id] = [bot]
+            else:
+                USER[user_id].append(bot)
+
+
+driver = get_driver()
+
+
+@driver.on_bot_connect
+@driver.on_bot_disconnect
+async def _(bot: Bot):
+    await refresh_bots()
+
+
+change_notice = on_notice(priority=1)
+
+
+@change_notice.handle()
+async def _(bot: Bot, event: FriendAddNoticeEvent):
+    await refresh_bots()
+
+
+# 01-06 16:56:51 [SUCCESS] nonebot | OneBot V11 **** | [notice.group_increase.approve]: {'time': 1672995411, 'self_id': ****, 'post_type': 'notice', 'notice_type': 'group_increase', 'sub_type': 'approve', 'user_id': ****, 'group_id': ****, 'operator_id': 0}
+# 01-06 16:58:09 [SUCCESS] nonebot | OneBot V11 **** | [notice.group_decrease.kick_me]: {'time': 1672995489, 'self_id': ****, 'post_type': 'notice', 'notice_type': 'group_decrease', 'sub_type': 'kick_me', 'user_id': ****, 'group_id': ****, 'operator_id': ****}
+@change_notice.handle()
+async def _(bot: Bot, event: GroupDecreaseNoticeEvent | GroupIncreaseNoticeEvent):
+    if bot.self_id == event.user_id:
+        await refresh_bots()
+
+
+def get_bot(user: User) -> Optional[Bot]:
+    """获取 Bot"""
+    bots = []
+    if user.user_type == "group":
+        bots = GROUP.get(user.user, [])
+
+    if user.user_type == "private":
+        bots = USER.get(user.user, [])
+
+    if not bots:
+        return
+
+    return random.choice(bots)
+
+
+async def get_groups() -> list[dict[str, Any]]:
+    """获取所有群号"""
+    all_groups: dict[int, dict[str, Any]] = {}
+    for bot in get_bots():
+        groups = await bot.get_group_list()
+        all_groups.update(
+            {
+                group["group_id"]: group
+                for group in groups
+                if group["group_id"] not in all_groups
+            }
+        )
+
+    return list(all_groups.values())