From 2c5593419806a11454f424b845655501c5316fc1 Mon Sep 17 00:00:00 2001
From: yu-xin-c <2182712990@qq.com>
Date: Sat, 29 Aug 2026 15:06:13 +0800
Subject: [PATCH] fix(agent): bound session message rate-limit state

---
 nanobot/agent/tools/session_messages.py   | 19 +++++--
 tests/tools/test_session_messages_tool.py | 60 +++++++++++++++++++++++
 2 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/nanobot/agent/tools/session_messages.py b/nanobot/agent/tools/session_messages.py
index 49a358de5..4be28f663 100644
--- a/nanobot/agent/tools/session_messages.py
+++ b/nanobot/agent/tools/session_messages.py
@@ -7,7 +7,7 @@ from __future__ import annotations
 import asyncio
 import json
 import time
-from collections import deque
+from collections import OrderedDict, deque
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any, Protocol
@@ -127,7 +127,7 @@ class SendSessionMessageTool(Tool):
         self._max_messages_per_minute = max_messages_per_minute
         self._schedule_later = schedule_later
         self._clock = clock or time.monotonic
-        self._sent_at: dict[str, deque[float]] = {}
+        self._sent_at: OrderedDict[str, deque[float]] = OrderedDict()
         self._pending_replies: dict[tuple[str, str], _PendingReply] = {}
         self._expiry_tasks: set[asyncio.Task[None]] = set()
         self._send_lock = asyncio.Lock()
@@ -240,8 +240,11 @@ class SendSessionMessageTool(Tool):
 
         async with self._send_lock:
             now = self._clock()
-            sent_at = self._sent_at.setdefault(source.session_key, deque())
             cutoff = now - _RATE_LIMIT_WINDOW_SECONDS
+            self._prune_expired_rate_limits(cutoff)
+            sent_at = self._sent_at.get(source.session_key)
+            if sent_at is None:
+                sent_at = deque[float]()
             while sent_at and sent_at[0] <= cutoff:
                 sent_at.popleft()
             if len(sent_at) >= self._max_messages_per_minute:
@@ -259,6 +262,8 @@ class SendSessionMessageTool(Tool):
                 input_role="user",
             ))
             sent_at.append(now)
+            self._sent_at[source.session_key] = sent_at
+            self._sent_at.move_to_end(source.session_key)
             self._cancel_pending_reply(reverse_wait_key)
             if timeout_seconds is not None:
                 self._cancel_pending_reply(wait_key)
@@ -271,6 +276,14 @@ class SendSessionMessageTool(Tool):
 
         return f"@{target.name}"
 
+    def _prune_expired_rate_limits(self, cutoff: float) -> None:
+        """Drop sources ordered by their most recent successful send."""
+        while self._sent_at:
+            _, sent_at = next(iter(self._sent_at.items()))
+            if sent_at[-1] > cutoff:
+                return
+            self._sent_at.popitem(last=False)
+
     @staticmethod
     def _validate_reply_timeout(
         expect_reply: bool,
diff --git a/tests/tools/test_session_messages_tool.py b/tests/tools/test_session_messages_tool.py
index f1970c4af..0397b8fc6 100644
--- a/tests/tools/test_session_messages_tool.py
+++ b/tests/tools/test_session_messages_tool.py
@@ -183,6 +183,66 @@ async def test_rate_limit_is_per_source_session_and_uses_a_rolling_minute(
     )
 
 
+@pytest.mark.asyncio
+async def test_rate_limit_releases_expired_source_state_and_keeps_recent_sources(
+    tmp_path: Path,
+) -> None:
+    sessions = SessionManager(tmp_path)
+    _persist(
+        sessions,
+        "websocket:a",
+        "websocket:b",
+        "websocket:c",
+        "websocket:target",
+    )
+    now = 0.0
+    tool = SendSessionMessageTool(
+        sessions=sessions,
+        bus=MessageBus(),
+        max_messages_per_minute=2,
+        clock=lambda: now,
+    )
+    target = _handle(sessions, "websocket:target").name
+
+    for source in ("websocket:a", "websocket:b"):
+        await tool.enqueue(
+            source_session_key=source,
+            target_handle=target,
+            content="initial",
+            expect_reply=False,
+        )
+    now = 30.0
+    await tool.enqueue(
+        source_session_key="websocket:a",
+        target_handle=target,
+        content="recent",
+        expect_reply=False,
+    )
+
+    now = 61.0
+    await tool.enqueue(
+        source_session_key="websocket:c",
+        target_handle=target,
+        content="trigger cleanup",
+        expect_reply=False,
+    )
+
+    assert set(tool._sent_at) == {"websocket:a", "websocket:c"}
+    await tool.enqueue(
+        source_session_key="websocket:a",
+        target_handle=target,
+        content="within rolling window",
+        expect_reply=False,
+    )
+    with pytest.raises(SessionMessageError, match="rate limit"):
+        await tool.enqueue(
+            source_session_key="websocket:a",
+            target_handle=target,
+            content="over limit",
+            expect_reply=False,
+        )
+
+
 @pytest.mark.asyncio
 async def test_reply_timeout_injects_a_user_input_back_into_the_source(
     tmp_path: Path,