From bbd7bbd7f593fcb2157e768ce09c830b209eac4a Mon Sep 17 00:00:00 2001
From: Xubin Ren <52506698+Re-bin@users.noreply.github.com>
Date: Fri, 19 Jun 2026 14:14:14 +0800
Subject: [PATCH] fix(mcp): avoid relying on progress notification root shape

---
 nanobot/agent/tools/mcp.py         | 26 ++++++++++++++++++++++----
 tests/agent/test_mcp_connection.py | 18 ++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py
index c34db90d..2fb7da89 100644
--- a/nanobot/agent/tools/mcp.py
+++ b/nanobot/agent/tools/mcp.py
@@ -47,12 +47,30 @@ _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":
+    payload = _mcp_jsonrpc_payload(message)
+    if _payload_value(payload, "method") != "notifications/progress":
         return False
 
-    params = getattr(root, "params", None)
-    return not isinstance(params, Mapping) or "progressToken" not in params
+    params = _payload_value(payload, "params")
+    return not _progress_params_have_token(params)
+
+
+def _mcp_jsonrpc_payload(message: Any) -> Any:
+    """Return the JSON-RPC payload across current and future MCP SDK shapes."""
+    envelope = getattr(message, "message", message)
+    return getattr(envelope, "root", None) or envelope
+
+
+def _payload_value(payload: Any, key: str) -> Any:
+    if isinstance(payload, Mapping):
+        return payload.get(key)
+    return getattr(payload, key, None)
+
+
+def _progress_params_have_token(params: Any) -> bool:
+    if isinstance(params, Mapping):
+        return "progressToken" in params
+    return hasattr(params, "progressToken") or hasattr(params, "progress_token")
 
 
 class _MalformedProgressNotificationFilter:
diff --git a/tests/agent/test_mcp_connection.py b/tests/agent/test_mcp_connection.py
index e145a71c..d70c6b37 100644
--- a/tests/agent/test_mcp_connection.py
+++ b/tests/agent/test_mcp_connection.py
@@ -36,6 +36,24 @@ def _mcp_notification(method: str, params: dict[str, Any] | None = None) -> Sess
     )
 
 
+def test_mcp_progress_detection_accepts_flattened_sdk_message_shape():
+    malformed = SimpleNamespace(
+        message=SimpleNamespace(
+            method="notifications/progress",
+            params={"progress": 20, "total": 600},
+        )
+    )
+    valid = SimpleNamespace(
+        message=SimpleNamespace(
+            method="notifications/progress",
+            params={"progressToken": "req-1", "progress": 25},
+        )
+    )
+
+    assert mcp_runtime._is_malformed_mcp_progress_notification(malformed) is True
+    assert mcp_runtime._is_malformed_mcp_progress_notification(valid) is False
+
+
 class _FakeMcpTool(Tool):
     def __init__(self, name: str) -> None:
         self._name = name