fix(notification): 限制实时消息队列积压并在 SSE 重连后一次推送 (#6862)

2026-11-01 04:47:03 +08:00 · 2026-10-02 21:32:18 +08:00
parent 5efdefe2a9
commit eb76fcafd6
4 changed files with 102 additions and 10 deletions
--- a/app/api/endpoints/system.py
+++ b/app/api/endpoints/system.py
@@ -679,8 +679,14 @@ async def get_message(
            while not runtime_stop_state.is_system_stopped:
                if await request.is_disconnected():
                    break
-                detail = message.get(role)
-                yield f"data: {detail or ''}\n\n"
+                # 一次推送全部积压消息，避免页面重新打开后按 3 秒一条逐条补弹
+                details = message.drain(role)
+                if details:
+                    for detail in details:
+                        yield f"data: {detail}\n\n"
+                else:
+                    # 空消息作为心跳，维持连接并定期检测客户端是否断开
+                    yield "data: \n\n"
                await asyncio.sleep(3)
        except asyncio.CancelledError:
            return
--- a/app/application/messaging/message.py
+++ b/app/application/messaging/message.py
@@ -36,6 +36,9 @@ _ALBUM_TRAILING_YEAR_RE = re.compile(
    r"(?:[\s\u3000]*[\(\[（【]\s*(?:19|20)\d{2}\s*[\)\]）】])+$"
 )
 _MESSAGE_QUEUE_STOP_TIMEOUT_SECONDS = 10.0
+# SSE 实时消息队列上限：前端关闭或切到后台时没有消费者，超出后丢弃最旧消息，
+# 避免长时间无人打开页面时队列无界增长，重新打开后又逐条补弹积压消息
+_SSE_QUEUE_MAXSIZE = 100


 class AsyncMessageQueryRepository(Protocol):
@@ -1200,11 +1203,12 @@ class MessageHelper(metaclass=Singleton):

    def __init__(self) -> None:
        """初始化系统消息队列和通知去重缓存。"""
-        self.sys_queue: queue.Queue[str] = queue.Queue()
+        self.sys_queue: queue.Queue[str] = queue.Queue(maxsize=_SSE_QUEUE_MAXSIZE)
        self._recent_notification_keys = TTLCache(region="message:notification", maxsize=500, ttl=60)

    def close(self) -> None:
-        """关闭通知去重缓存；真实收敛后由生命周期释放单例身份。"""
+        """清空未推送的实时消息并关闭通知去重缓存；真实收敛后由生命周期释放单例身份。"""
+        self.drain()
        self._recent_notification_keys.close()

    @staticmethod
@@ -1264,7 +1268,7 @@ class MessageHelper(metaclass=Singleton):
            title = "插件通知"
        if self._is_recent_system_notification(message, role, title=title, note=note):
            return
-        self.sys_queue.put(json.dumps({
+        self._put_drop_oldest(json.dumps({
            "type": role,
            "title": title,
            "text": message,
@@ -1274,14 +1278,48 @@ class MessageHelper(metaclass=Singleton):
            "level": level,
        }))

+    def _put_drop_oldest(self, payload: str) -> None:
+        """
+        入队实时消息，队列已满时丢弃最旧的一条再重试。
+
+        多个生产线程并发写入时，腾出的位置可能先被其他生产者占用而再次遇到队列已满；
+        每次重试前都有其他生产者写入成功，队列上限远大于并发生产者数，很快就能写入。
+        """
+        while True:
+            try:
+                self.sys_queue.put_nowait(payload)
+                return
+            except queue.Full:
+                try:
+                    self.sys_queue.get_nowait()
+                    logger.debug("实时消息队列已满，丢弃最旧的一条消息")
+                except queue.Empty:
+                    pass
+
    def get(self, role: str = "system") -> Optional[str]:
        """
        取消息
        :param role: 兼容旧参数，当前所有 SSE 消息共用一个队列
        """
-        if not self.sys_queue.empty():
-            return self.sys_queue.get(block=False)
-        return None
+        try:
+            return self.sys_queue.get_nowait()
+        except queue.Empty:
+            return None
+
+    def drain(self, role: str = "system") -> list[str]:
+        """
+        按入队顺序取出当前积压的全部实时消息。
+
+        最多取出队列上限条，生产方持续写入时也能在有限次数内返回。
+        :param role: 兼容旧参数，当前所有 SSE 消息共用一个队列
+        """
+        messages: list[str] = []
+        for _ in range(_SSE_QUEUE_MAXSIZE):
+            detail = self.get(role)
+            if detail is None:
+                break
+            messages.append(detail)
+        return messages


 def stop_message(
--- a/tests/test_message_notifications.py
+++ b/tests/test_message_notifications.py
@@ -3,6 +3,7 @@ import json
 from unittest.mock import Mock

 from app.api.endpoints.message import clear_notification_message, get_notification_message
+from app.api.endpoints.system import get_message
 from app.chain.base import ChainBase
 from app.domain.context import Context, MediaInfo, TorrentInfo
 from app.domain.meta.metabase import MetaBase
@@ -10,7 +11,7 @@ from app.db.session import AsyncSessionFactory, SessionFactory
 from app.db.oper.message import MessageOper
 from app.db.models.message import Message as MessageModel
 from app.db.oper.systemconfig import SystemConfigOper
-from app.application.messaging.message import MessageHelper, MessageQueryService
+from app.application.messaging.message import _SSE_QUEUE_MAXSIZE, MessageHelper, MessageQueryService
 from app.schemas.message import Message, MessageClearScope
 from app.schemas.types import MediaType, MessageType, SystemConfigKey

@@ -166,6 +167,51 @@ def test_plugin_helper_message_deduplicates_recent_sse_messages() -> None:
    assert helper.get() is None


+def test_helper_message_queue_drops_oldest_when_full() -> None:
+    """
+    没有 SSE 消费者时实时消息队列有上限，超出后丢弃最旧消息，保留最新消息。
+    """
+    helper = MessageHelper()
+    _reset_message_helper(helper)
+
+    for index in range(_SSE_QUEUE_MAXSIZE + 5):
+        helper.put(f"积压消息 {index}", role="system", title="队列上限")
+
+    assert helper.sys_queue.qsize() == _SSE_QUEUE_MAXSIZE
+    texts = [json.loads(detail)["text"] for detail in helper.drain()]
+    assert texts == [f"积压消息 {index}" for index in range(5, _SSE_QUEUE_MAXSIZE + 5)]
+    assert helper.get() is None
+
+
+def test_sse_endpoint_flushes_backlog_in_one_tick() -> None:
+    """
+    SSE 重新连接后应一次推送全部积压消息，而不是每 3 秒推送一条。
+    """
+    helper = MessageHelper()
+    _reset_message_helper(helper)
+    for index in range(3):
+        helper.put(f"重连积压 {index}", role="plugin", title="积压推送")
+
+    class _ConnectedRequest:
+        async def is_disconnected(self) -> bool:
+            return False
+
+    async def _read_backlog() -> list[str]:
+        response = await get_message(request=_ConnectedRequest(), role="notification", _=None)
+        stream = response.body_iterator
+        try:
+            # 3 秒间隔只发生在一轮推送之后，整轮积压应在超时前全部到达
+            return [await asyncio.wait_for(anext(stream), timeout=1) for _ in range(3)]
+        finally:
+            await stream.aclose()
+
+    frames = asyncio.run(_read_backlog())
+
+    texts = [json.loads(frame.removeprefix("data: ").strip())["text"] for frame in frames]
+    assert texts == [f"重连积压 {index}" for index in range(3)]
+    assert helper.get() is None
+
+
 def test_agent_helper_message_does_not_enter_sse_queue() -> None:
    """
    智能体消息不进入前端 SSE 队列。
--- a/tests/test_message_queue_shutdown.py
+++ b/tests/test_message_queue_shutdown.py
@@ -107,15 +107,17 @@ def test_bound_clients_keep_distinct_chain_callbacks_without_thread(monkeypatch)


 def test_stop_message_closes_and_releases_explicit_message_helper(monkeypatch) -> None:
-    """消息通知缓存收敛后必须释放单例，避免下一轮复用已关闭缓存。"""
+    """消息通知缓存收敛后必须释放单例，避免下一轮复用已关闭缓存；未推送的实时消息一并清空。"""
    monkeypatch.setattr(Singleton, "_instances", {})
    helper = MessageHelper()
    close = MagicMock()
    monkeypatch.setattr(helper._recent_notification_keys, "close", close)
+    helper.sys_queue.put_nowait("未推送的实时消息")

    assert stop_message(message_helper=helper) is True

    close.assert_called_once_with()
+    assert helper.sys_queue.empty()
    assert MessageHelper.get_existing_instance() is None
    assert MessageHelper() is not helper