From d33b06f44d9716ce0ff4d40c8bdfaaa42ccfcffb Mon Sep 17 00:00:00 2001
From: yunyancuo <3468440670@qq.com>
Date: Wed, 22 Jul 2026 15:32:56 +0800
Subject: [PATCH 1/2] fix: move now = datetime.now() inside the lock to avoid
 stale timestamp in concurrent scenario
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When multiple coroutines wait for the same session lock, they all capture

ow = datetime.now() before entering the lock. By the time a coroutine
actually acquires the lock, the captured timestamp is stale — it reflects
the time before the wait, not the actual acquisition time. This causes
expired timestamps to not be cleaned and the rate limit window to be
calculated from the wrong reference time.

Moving
ow = datetime.now() inside the lock ensures every coroutine
uses the current time at the moment it acquires the lock, fixing the
stale-timestamp bug for the stall strategy.

Co-authored-by: yunyancuo <3468440670@qq.com>
---
 .../core/pipeline/rate_limit_check/stage.py   |  2 +-
 tests/test_rate_limit_stage.py                | 64 +++++++++++++++++++
 2 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_rate_limit_stage.py

diff --git a/astrbot/core/pipeline/rate_limit_check/stage.py b/astrbot/core/pipeline/rate_limit_check/stage.py
index f5d12fdfa7..ab9f68ddad 100644
--- a/astrbot/core/pipeline/rate_limit_check/stage.py
+++ b/astrbot/core/pipeline/rate_limit_check/stage.py
@@ -55,9 +55,9 @@ async def process(
 
         """
         session_id = event.session_id
-        now = datetime.now()
 
         async with self.locks[session_id]:  # 确保同一会话不会并发修改队列
+            now = datetime.now()
             # 检查并处理限流，可能需要多次检查直到满足条件
             while True:
                 timestamps = self.event_timestamps[session_id]
diff --git a/tests/test_rate_limit_stage.py b/tests/test_rate_limit_stage.py
new file mode 100644
index 0000000000..03c238fd21
--- /dev/null
+++ b/tests/test_rate_limit_stage.py
@@ -0,0 +1,64 @@
+import asyncio
+from datetime import datetime as real_datetime
+from datetime import timedelta
+
+import pytest
+
+from astrbot.core.pipeline.rate_limit_check import stage as rate_limit_stage
+
+
+class FakeEvent:
+    """Minimal message event used by the rate-limit stage tests."""
+
+    session_id = "test-session"
+
+    def stop_event(self) -> None:
+        """Stop event propagation for discard-strategy compatibility."""
+
+
+@pytest.mark.asyncio
+async def test_stalled_concurrent_events_use_current_time_after_lock(monkeypatch):
+    """Ensure queued events do not reuse timestamps captured before lock waits."""
+    virtual_seconds = 0.0
+    sleep_durations: list[float] = []
+    real_sleep = asyncio.sleep
+    base_time = real_datetime(2026, 1, 1)
+
+    class FakeDateTime:
+        """Return time from the deterministic virtual clock."""
+
+        @classmethod
+        def now(cls) -> real_datetime:
+            """Return the current virtual wall-clock time.
+
+            Returns:
+                Current virtual time.
+            """
+            return base_time + timedelta(seconds=virtual_seconds)
+
+    async def fake_sleep(duration: float) -> None:
+        """Advance virtual time after allowing concurrent tasks to queue.
+
+        Args:
+            duration: Requested sleep duration in seconds.
+        """
+        nonlocal virtual_seconds
+        sleep_durations.append(duration)
+        target_time = virtual_seconds + duration
+        await real_sleep(0)
+        virtual_seconds = target_time
+
+    monkeypatch.setattr(rate_limit_stage, "datetime", FakeDateTime)
+    monkeypatch.setattr(rate_limit_stage.asyncio, "sleep", fake_sleep)
+    monkeypatch.setattr(rate_limit_stage.logger, "info", lambda *args, **kwargs: None)
+
+    limiter = rate_limit_stage.RateLimitStage()
+    limiter.rate_limit_count = 2
+    limiter.rate_limit_time = timedelta(seconds=60)
+    limiter.rl_strategy = "stall"
+
+    await asyncio.gather(*(limiter.process(FakeEvent()) for _ in range(5)))
+
+    assert sleep_durations == pytest.approx([60.3, 60.3])
+    timestamps = list(limiter.event_timestamps[FakeEvent.session_id])
+    assert timestamps == sorted(timestamps)

From cbc6b0b6515eafc2640fd98f64ee586cd45710a6 Mon Sep 17 00:00:00 2001
From: yunyancuo <3468440670@qq.com>
Date: Wed, 22 Jul 2026 15:42:50 +0800
Subject: [PATCH 2/2] test: address review - subclass datetime, derive expected
 stall

---
 tests/test_rate_limit_stage.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tests/test_rate_limit_stage.py b/tests/test_rate_limit_stage.py
index 03c238fd21..b7ff1cb789 100644
--- a/tests/test_rate_limit_stage.py
+++ b/tests/test_rate_limit_stage.py
@@ -24,8 +24,8 @@ async def test_stalled_concurrent_events_use_current_time_after_lock(monkeypatch
     real_sleep = asyncio.sleep
     base_time = real_datetime(2026, 1, 1)
 
-    class FakeDateTime:
-        """Return time from the deterministic virtual clock."""
+    class FakeDateTime(real_datetime):
+        """Subclass of datetime with a deterministic now()."""
 
         @classmethod
         def now(cls) -> real_datetime:
@@ -59,6 +59,8 @@ async def fake_sleep(duration: float) -> None:
 
     await asyncio.gather(*(limiter.process(FakeEvent()) for _ in range(5)))
 
-    assert sleep_durations == pytest.approx([60.3, 60.3])
+    margin = 0.3
+    expected_stall = limiter.rate_limit_time.total_seconds() + margin
+    assert sleep_durations == pytest.approx([expected_stall, expected_stall])
     timestamps = list(limiter.event_timestamps[FakeEvent.session_id])
     assert timestamps == sorted(timestamps)