实现current_rate_limit_token (#23)

* 实现current_rate_limit_token

* 升级版本号

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nonebot-plugin-access-control"
-version = "1.1.3"
+version = "1.1.4"
 description = ""
 authors = ["ssttkkl <[email protected]>"]
 license = "MIT"
@@ -13,7 +13,7 @@ packages = [
 [tool.poetry.dependencies]
 python = "^3.9"
 nonebot2 = "^2.1.0"
-nonebot-plugin-access-control-api = "^1.1.1"
+nonebot-plugin-access-control-api = "^1.1.2"
 nonebot-plugin-apscheduler = ">=0.3.0"
 nonebot-plugin-session = "^0.2.0"
 nonebot-plugin-orm = ">=0.5.0, <1.0.0"

src/nonebot_plugin_ac_demo/matcher_demo.py

@@ -1,5 +1,8 @@
 from nonebot import on_command
 from nonebot.internal.matcher import Matcher
+from nonebot_plugin_access_control_api.service.contextvars import (
+    current_rate_limit_token,
+)
 
 from .plugin_service import plugin_service
 
@@ -33,3 +36,21 @@ async def _(matcher: Matcher):
 @c_service.patch_handler()
 async def _(matcher: Matcher):
     await matcher.send("c")
+
+
+d_counter = 0
+
+d_matcher = on_command("d", priority=99)
+d_service = plugin_service.create_subservice("d")
+
+
+@d_matcher.handle()
+@d_service.patch_handler()
+async def _(matcher: Matcher):
+    global d_counter
+    d_counter += 1
+    if d_counter % 2 == 0:
+        await current_rate_limit_token.get().retire()
+        await matcher.send("retired")
+    else:
+        await matcher.send("d")

src/nonebot_plugin_access_control/service/_impl/patcher.py

@@ -9,6 +9,9 @@
 from nonebot_plugin_access_control_api.service import get_nonebot_service
 from nonebot_plugin_access_control_api.models.rate_limit import AcquireTokenResult
 from nonebot_plugin_access_control_api.service.interface.patcher import IServicePatcher
+from nonebot_plugin_access_control_api.service.contextvars import (
+    current_rate_limit_token,
+)
 from nonebot_plugin_access_control_api.errors import (
     RateLimitedError,
     PermissionDeniedError,
@@ -76,22 +79,23 @@ async def wrapped_func(*args, **kwargs):
                     await self.handle_rate_limited(matcher, result)
                     return
 
-                matcher.state["ac_token"] = result.token
-
+                t = current_rate_limit_token.set(result.token)
                 try:
                     return await func(*args, **kwargs)
                 except BaseException as e:
                     if retire_on_throw:
                         await result.token.retire()
                     raise e
+                finally:
+                    current_rate_limit_token.reset(t)
 
             return wrapped_func
 
         return decorator
 
 
 @run_preprocessor
-async def check(matcher: Matcher, bot: Bot, event: Event):
+async def check(bot: Bot, event: Event, matcher: Matcher):
     service = ServicePatcherImpl._matcher_service_mapping.get(type(matcher), None)
     if service is None:
         return

src/tests/test_rate_limit.py

@@ -111,3 +111,35 @@ async def test_rate_limit_overwrite(app: App):
         event = fake_ob11_group_message_event("/c")
         ctx.receive_event(bot, event)
         ctx.should_call_send(event, "c")
+
+
+@pytest.mark.asyncio
+async def test_rate_limit_retire(app: App):
+    from nonebot.adapters.onebot.v11 import Bot
+    from nonebot_plugin_access_control_api.service import get_nonebot_service
+
+    from nonebot_plugin_ac_demo.matcher_demo import d_matcher
+
+    # service: nonebot
+    # subject: all
+    # span: 1s
+    # limit: 2
+    await get_nonebot_service().add_rate_limit_rule("all", timedelta(seconds=1), 2)
+
+    async with app.test_matcher(d_matcher) as ctx:
+        bot = ctx.create_bot(base=Bot, self_id=str(SELF_ID))
+        event = fake_ob11_group_message_event("/d")
+        ctx.receive_event(bot, event)
+        ctx.should_call_send(event, "d")
+
+    async with app.test_matcher(d_matcher) as ctx:
+        bot = ctx.create_bot(base=Bot, self_id=str(SELF_ID))
+        event = fake_ob11_group_message_event("/d")
+        ctx.receive_event(bot, event)
+        ctx.should_call_send(event, "retired")
+
+    async with app.test_matcher(d_matcher) as ctx:
+        bot = ctx.create_bot(base=Bot, self_id=str(SELF_ID))
+        event = fake_ob11_group_message_event("/d")
+        ctx.receive_event(bot, event)
+        ctx.should_call_send(event, "d")

src/tests/test_service_handler.py

@@ -17,6 +17,7 @@ async def test_service_ls_handler(app: App):
         "│   ├── a\n"
         "│   └── b\n"
         "├── c\n"
+        "├── d\n"
         "└── tick"
     )