👷 fix by ruff

nonebot_bison/platform/bilibili/scheduler.py

@@ -1,14 +1,17 @@
+from collections.abc import Callable
 from datetime import datetime, timedelta
 import json
 import random
 from typing import TYPE_CHECKING, ClassVar, TypeVar
 from typing_extensions import override
 
-from httpx import AsyncClient
+from httpx import AsyncClient, Response
 from nonebot import logger, require
 from playwright.async_api import Cookie
 
+from nonebot_bison.config import config
 from nonebot_bison.config.db_model import Cookie as CookieModel
+from nonebot_bison.config.db_model import Target
 from nonebot_bison.utils import Site, http_client
 from nonebot_bison.utils.site import CookieClientManager
 
@@ -23,6 +26,8 @@ B = TypeVar("B", bound="Bilibili")
 
 class BilibiliClientManager(CookieClientManager):
     _default_cookie_cd = timedelta(seconds=120)
+    _current_user_cookie: CookieModel | None = None
+    _site_name = "bilibili.com"
 
     async def _get_cookies(self) -> list[Cookie]:
         browser = await get_browser()
@@ -57,6 +62,40 @@ class BilibiliClientManager(CookieClientManager):
         )
         return cookie
 
+    def _generate_hook(self, cookie: CookieModel) -> Callable:
+        """hook 函数生成器，用于回写请求状态到数据库"""
+
+        async def _response_hook(resp: Response):
+            await resp.aread()
+            if resp.status_code == 200 and "-352" not in resp.text:
+                logger.trace(f"请求成功: {cookie.id} {resp.request.url}")
+                cookie.status = "success"
+            else:
+                logger.warning(f"请求失败: {cookie.id} {resp.request.url}, 状态码: {resp.status_code}")
+                cookie.status = "failed"
+                self._current_user_cookie = None
+            cookie.last_usage = datetime.now()
+            await config.update_cookie(cookie)
+
+        return _response_hook
+
+    async def _get_next_user_cookie(self) -> CookieModel | None:
+        """选择下一个用户 cookie"""
+        cookies = await config.get_cookie(self._site_name, is_anonymous=False)
+        available_cookies = [cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now()]
+        if not available_cookies:
+            return None
+        cookie = min(available_cookies, key=lambda x: x.last_usage)
+        return cookie
+
+    async def _choose_cookie(self, target: Target | None) -> CookieModel:
+        """选择 cookie 的具体算法"""
+        if self._current_user_cookie is None:
+            self._current_user_cookie = await self._get_next_user_cookie()
+        if self._current_user_cookie:
+            return self._current_user_cookie
+        return (await config.get_cookie(self._site_name, is_anonymous=True))[0]
+
     @override
     async def refresh_client(self):
         await self._refresh_anonymous_cookie()

nonebot_bison/utils/site.py

@@ -124,8 +124,8 @@ class CookieClientManager(ClientManager):
     async def _choose_cookie(self, target: Target | None) -> Cookie:
         """选择 cookie 的具体算法"""
         cookies = await config.get_cookie(self._site_name, target)
-        avaliable_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
-        cookie = min(avaliable_cookies, key=lambda x: x.last_usage)
+        available_cookies = (cookie for cookie in cookies if cookie.last_usage + cookie.cd < datetime.now())
+        cookie = min(available_cookies, key=lambda x: x.last_usage)
         return cookie
 
     async def get_client(self, target: Target | None) -> AsyncClient: