From eb76fcafd66d46220b586653ba8ad46dea92c822 Mon Sep 17 00:00:00 2001 From: Cat <86354789+MiCat-S@users.noreply.github.com> Date: Fri, 2 Oct 2026 21:32:18 +0800 Subject: [PATCH] =?UTF-8?q?fix(notification):=20=E9=99=90=E5=88=B6?= =?UTF-8?q?=E5=AE=9E=E6=97=B6=E6=B6=88=E6=81=AF=E9=98=9F=E5=88=97=E7=A7=AF?= =?UTF-8?q?=E5=8E=8B=E5=B9=B6=E5=9C=A8=20SSE=20=E9=87=8D=E8=BF=9E=E5=90=8E?= =?UTF-8?q?=E4=B8=80=E6=AC=A1=E6=8E=A8=E9=80=81=20(#6862)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/api/endpoints/system.py | 10 ++++-- app/application/messaging/message.py | 50 ++++++++++++++++++++++++---- tests/test_message_notifications.py | 48 +++++++++++++++++++++++++- tests/test_message_queue_shutdown.py | 4 ++- 4 files changed, 102 insertions(+), 10 deletions(-) diff --git a/app/api/endpoints/system.py b/app/api/endpoints/system.py index 493b4cae0..861f6c544 100644 --- 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 diff --git a/app/application/messaging/message.py b/app/application/messaging/message.py index 235bb9853..71e69c467 100644 --- 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( diff --git a/tests/test_message_notifications.py b/tests/test_message_notifications.py index 05916d9fc..9143e7209 100644 --- 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 队列。 diff --git a/tests/test_message_queue_shutdown.py b/tests/test_message_queue_shutdown.py index 944f24373..4cea65850 100644 --- 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