From f9511049c498c015831d6ff0baa8381ea1625a8c Mon Sep 17 00:00:00 2001
From: yu-xin-c <2182712990@qq.com>
Date: Wed, 17 Jun 2026 01:09:08 +0800
Subject: [PATCH] fix(mcp): ignore malformed progress notifications

---
 nanobot/agent/tools/mcp.py         | 53 ++++++++++++++++++++++++++++++
 tests/agent/test_mcp_connection.py | 41 +++++++++++++++++++++++
 2 files changed, 94 insertions(+)

diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py
index 181c4e9f..c34db90d 100644
--- a/nanobot/agent/tools/mcp.py
+++ b/nanobot/agent/tools/mcp.py
@@ -46,6 +46,58 @@ _RELOAD_LOCKS: WeakKeyDictionary[Any, asyncio.Lock] = WeakKeyDictionary()
 _ReconnectCallback = Callable[[str, str, Tool], Awaitable[Tool | None]]
 
 
+def _is_malformed_mcp_progress_notification(message: Any) -> bool:
+    root = getattr(getattr(message, "message", None), "root", None)
+    if getattr(root, "method", None) != "notifications/progress":
+        return False
+
+    params = getattr(root, "params", None)
+    return not isinstance(params, Mapping) or "progressToken" not in params
+
+
+class _MalformedProgressNotificationFilter:
+    def __init__(self, read_stream: Any, server_name: str) -> None:
+        self._read_stream = read_stream
+        self._server_name = server_name
+        self._iterator: Any | None = None
+
+    async def __aenter__(self) -> "_MalformedProgressNotificationFilter":
+        await self._read_stream.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> Any:
+        return await self._read_stream.__aexit__(exc_type, exc, tb)
+
+    def __aiter__(self) -> "_MalformedProgressNotificationFilter":
+        self._iterator = self._read_stream.__aiter__()
+        return self
+
+    async def __anext__(self) -> Any:
+        if self._iterator is None:
+            self._iterator = self._read_stream.__aiter__()
+
+        while True:
+            message = await self._iterator.__anext__()
+            if _is_malformed_mcp_progress_notification(message):
+                logger.debug(
+                    "MCP server '{}': dropped progress notification without progressToken",
+                    self._server_name,
+                )
+                continue
+            return message
+
+    async def aclose(self) -> None:
+        close = getattr(self._read_stream, "aclose", None)
+        if close is not None:
+            await close()
+
+
+def _filter_malformed_mcp_progress_notifications(read_stream: Any, server_name: str) -> Any:
+    if not all(hasattr(read_stream, name) for name in ("__aenter__", "__aexit__", "__aiter__")):
+        return read_stream
+    return _MalformedProgressNotificationFilter(read_stream, server_name)
+
+
 def _sanitize_name(name: str) -> str:
     """Sanitize an MCP-derived name for model API compatibility."""
     return _SANITIZE_RE.sub("_", re.sub(r"[^a-zA-Z0-9_-]", "_", name))
@@ -681,6 +733,7 @@ async def connect_mcp_servers(
                 await server_stack.aclose()
                 return name, None
 
+            read = _filter_malformed_mcp_progress_notifications(read, name)
             session = await server_stack.enter_async_context(ClientSession(read, write))
             await session.initialize()
 
diff --git a/tests/agent/test_mcp_connection.py b/tests/agent/test_mcp_connection.py
index d9b32520..e145a71c 100644
--- a/tests/agent/test_mcp_connection.py
+++ b/tests/agent/test_mcp_connection.py
@@ -8,9 +8,11 @@ from types import SimpleNamespace
 from typing import Any
 from unittest.mock import MagicMock
 
+import anyio
 import pytest
 from mcp import types as mcp_types
 from mcp.shared.exceptions import McpError
+from mcp.shared.message import SessionMessage
 from mcp.types import ErrorData
 
 from nanobot.agent.loop import AgentLoop
@@ -22,6 +24,18 @@ from nanobot.config.loader import load_config, save_config
 from nanobot.config.schema import MCPServerConfig
 
 
+def _mcp_notification(method: str, params: dict[str, Any] | None = None) -> SessionMessage:
+    return SessionMessage(
+        message=mcp_types.JSONRPCMessage(
+            mcp_types.JSONRPCNotification(
+                jsonrpc="2.0",
+                method=method,
+                params=params,
+            )
+        )
+    )
+
+
 class _FakeMcpTool(Tool):
     def __init__(self, name: str) -> None:
         self._name = name
@@ -56,6 +70,33 @@ def _make_loop(tmp_path, *, mcp_servers: dict | None = None) -> AgentLoop:
     )
 
 
+@pytest.mark.asyncio
+async def test_mcp_read_filter_drops_progress_notifications_without_progress_token():
+    send, receive = anyio.create_memory_object_stream(4)
+    malformed_progress = _mcp_notification(
+        "notifications/progress",
+        {"progress": 20, "total": 600, "message": "Polling"},
+    )
+    tool_change = _mcp_notification("notifications/tools/list_changed")
+    valid_progress = _mcp_notification(
+        "notifications/progress",
+        {"progressToken": "req-1", "progress": 25, "total": 600, "message": "Polling"},
+    )
+
+    await send.send(malformed_progress)
+    await send.send(tool_change)
+    await send.send(valid_progress)
+    await send.aclose()
+
+    wrapped = mcp_runtime._filter_malformed_mcp_progress_notifications(receive, "brightdata")
+    forwarded = []
+    async with wrapped:
+        async for message in wrapped:
+            forwarded.append(message)
+
+    assert forwarded == [tool_change, valid_progress]
+
+
 @pytest.mark.asyncio
 async def test_connect_mcp_retries_when_no_servers_connect(tmp_path, monkeypatch: pytest.MonkeyPatch):
     loop = _make_loop(tmp_path)