✨ 匿名cookie和用户cookie一起调度

nonebot_bison/config/db_config.py

@@ -272,7 +272,20 @@ class DBConfig:
                 res = [cookie for cookie in res if cookie.id in ids]
             return res
 
-    async def add_cookie(self, platform_name: str, content: str) -> int:
+    async def get_unviersal_cookie(self, platform_name: str = None, target: T_Target = None) -> list[Cookie]:
+        async with create_session() as sess:
+            query = select(Cookie).distinct().where(Cookie.is_universal == True)  # noqa: E712
+            if platform_name:
+                query = query.where(Cookie.platform_name == platform_name)
+            query = query.outerjoin(CookieTarget).options(selectinload(Cookie.targets))
+            res = (await sess.scalars(query)).all()
+            if target:
+                query = select(CookieTarget.cookie_id).join(Target).where(Target.target == target)
+                ids = set((await sess.scalars(query)).all())
+                res = [cookie for cookie in res if cookie.id in ids]
+            return res
+
+    async def add_cookie_with_content(self, platform_name: str, content: str) -> int:
         async with create_session() as sess:
             cookie = Cookie(platform_name=platform_name, content=content)
             sess.add(cookie)
@@ -280,6 +293,13 @@ class DBConfig:
             await sess.refresh(cookie)
             return cookie.id
 
+    async def add_cookie(self, cookie: Cookie) -> int:
+        async with create_session() as sess:
+            sess.add(cookie)
+            await sess.commit()
+            await sess.refresh(cookie)
+            return cookie.id
+
     async def update_cookie(self, cookie: Cookie):
         async with create_session() as sess:
             cookie_in_db: Cookie | None = await sess.scalar(select(Cookie).where(Cookie.id == cookie.id))
@@ -306,6 +326,15 @@ class DBConfig:
             )
             return list((await sess.scalars(query)).all())
 
+    async def get_universal_cookie(self, platform_name: str) -> list[Cookie]:
+        async with create_session() as sess:
+            query = (
+                select(Cookie)
+                .where(Cookie.platform_name == platform_name)
+                .where(Cookie.is_universal == True)  # noqa: E712
+            )
+            return list((await sess.scalars(query)).all())
+
     async def add_cookie_target(self, target: T_Target, platform_name: str, cookie_id: int):
         async with create_session() as sess:
             target_obj = await sess.scalar(

nonebot_bison/config/db_model.py

@@ -74,10 +74,15 @@ class Cookie(Model):
     id: Mapped[int] = mapped_column(primary_key=True)
     platform_name: Mapped[str] = mapped_column(String(20))
     content: Mapped[str] = mapped_column(String(1024))
+    # 最后使用的时刻
     last_usage: Mapped[datetime.datetime] = mapped_column(DateTime, default=datetime.datetime(1970, 1, 1))
+    # Cookie 当前的状态
     status: Mapped[str] = mapped_column(String(20), default="")
+    # 使用一次之后，需要的冷却时间
     cd: Mapped[int] = mapped_column(default=0)
+    # 是否是通用 Cookie，默认用于匿名 Cookie
     is_universal: Mapped[bool] = mapped_column(default=False)
+    # 标签，扩展用
     tags: Mapped[dict[str, Any]] = mapped_column(JSON().with_variant(JSONB, "postgresql"), default={})
 
     targets: Mapped[list["CookieTarget"]] = relationship(back_populates="cookie")

nonebot_bison/scheduler/manager.py

@@ -30,6 +30,8 @@ async def init_scheduler():
         else:
             _schedule_class_platform_dict[site].append(platform_name)
     for site, target_list in _schedule_class_dict.items():
+        if hasattr(site.client_mgr, "_cookie_client_manger_"):
+            await site.client_mgr.init_universal_cookie()
         if not plugin_config.bison_use_browser and site.require_browser:
             logger.warning(f"{site.name} requires browser, it will not schedule.")
             continue

nonebot_bison/utils/site.py

@@ -42,6 +42,7 @@ class DefaultClientManager(ClientManager):
 
 
 class CookieClientManager(ClientManager):
+    _cookie_client_manger_ = True
     _platform_name: str
     _cookie_cd: int = 10
 
@@ -49,9 +50,6 @@ class CookieClientManager(ClientManager):
         """hook函数生成器，用于回写请求状态到数据库"""
 
         async def _response_hook(resp: httpx.Response):
-            if cookie.content == "{}":
-                logger.debug(f"未携带bison cookie: {resp.request.url}")
-                return
             if resp.status_code == 200:
                 logger.debug(f"请求成功: {cookie.id} {resp.request.url}")
                 cookie.status = "success"
@@ -63,18 +61,30 @@ class CookieClientManager(ClientManager):
 
         return _response_hook
 
+    @classmethod
+    async def init_universal_cookie(cls):
+        """移除已有的匿名cookie，添加一个新的匿名cookie"""
+        universal_cookies = await config.get_unviersal_cookie(cls._platform_name)
+        for cookie in universal_cookies:
+            await config.delete_cookie(cookie.id)
+        universal_cookie = Cookie(platform_name=cls._platform_name, content="{}", is_universal=True)
+        await config.add_cookie(universal_cookie)
+
+    @classmethod
+    async def init_cookie(cls, cookie: Cookie):
+        """初始化cookie，添加用户cookie时使用"""
+        cookie.cd = cls._cookie_cd
+
     async def _choose_cookie(self, target: Target) -> Cookie:
         """选择 cookie 的具体算法"""
         if not target:
             return Cookie(content="{}")
-        cookies = await config.get_cookie_by_target(target, self._platform_name)
-        cookies = (
-            cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=self._cookie_cd) < datetime.now()
-        )
-        cookies = list(cookies)
-        if not cookies:
-            return Cookie(content="{}")
-        cookie = max(cookies, key=lambda x: x.last_usage)
+        cookies = [
+            *(await config.get_cookie_by_target(target, self._platform_name)),
+            *(await config.get_universal_cookie(self._platform_name)),
+        ]
+        cookies = (cookie for cookie in cookies if cookie.last_usage + timedelta(seconds=cookie.cd) < datetime.now())
+        cookie = min(cookies, key=lambda x: x.last_usage)
         return cookie
 
     async def _check_cookie(self, cookie: Cookie) -> Cookie:
@@ -85,10 +95,10 @@ class CookieClientManager(ClientManager):
         """获取 client，根据 target 选择 cookie"""
         client = http_client()
         cookie = await self._choose_cookie(target)
-        if cookie.content != "{}":
-            logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}")
+        if cookie.is_universal:
+            logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用匿名cookie")
         else:
-            logger.debug(f"平台 {self._platform_name} 未获取到用户cookie, 使用空cookie")
+            logger.debug(f"平台 {self._platform_name} 获取到用户cookie: {cookie.id}")
 
         return await self._assemble_client(client, cookie)
 

tests/config/test_cookie.py

@@ -50,7 +50,7 @@ async def test_cookie_by_user(app: App, init_scheduler):
         tags=[],
     )
 
-    await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie")
+    await config.add_cookie_with_content(TargetQQGroup(group_id=123), "weibo", "cookie")
 
     cookies = await config.get_cookie(TargetQQGroup(group_id=123))
     cookie = cookies[0]
@@ -93,7 +93,7 @@ async def test_cookie_target_by_target(app: App, init_scheduler):
         tags=[],
     )
 
-    id = await config.add_cookie(TargetQQGroup(group_id=123), "weibo", "cookie")
+    id = await config.add_cookie_with_content(TargetQQGroup(group_id=123), "weibo", "cookie")
 
     await config.add_cookie_target(T_Target("weibo_id"), "weibo", id)