From c2447a6cb9d4dfd29a51eeb0dfb353c77746f1d0 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sun, 5 Jul 2026 22:30:04 -0400
Subject: [PATCH] Route notifications/tasks/status via SDK v2 notification
 binding (custom notifs bypass message_handler)

---
 fastmcp_slim/fastmcp/client/client.py         | 44 +++++++++++++++++--
 .../fastmcp/client/transports/base.py         |  6 ++-
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/fastmcp_slim/fastmcp/client/client.py b/fastmcp_slim/fastmcp/client/client.py
index 0458ee50e..aab21fed7 100644
--- a/fastmcp_slim/fastmcp/client/client.py
+++ b/fastmcp_slim/fastmcp/client/client.py
@@ -17,7 +17,12 @@ import httpx
 import mcp_types
 from exceptiongroup import catch
 from mcp import ClientSession, MCPError
-from mcp_types import GetTaskResult, TaskStatusNotification
+from mcp.client.extension import NotificationBinding
+from mcp_types import (
+    GetTaskResult,
+    TaskStatusNotification,
+    TaskStatusNotificationParams,
+)
 from pydantic import AnyUrl
 
 import fastmcp as fastmcp
@@ -319,6 +324,10 @@ class Client(
             "message_handler": message_handler or TaskNotificationHandler(self),
             "read_timeout_seconds": read_timeout_seconds,
             "client_info": client_info,
+            # SDK v2 does not carry `notifications/tasks/status` in any protocol
+            # version's core notification tables, so it is never tee'd to the
+            # message_handler; a binding routes it to Task objects instead.
+            "notification_bindings": [self._task_status_binding()],
         }
 
         if roots is not None:
@@ -455,6 +464,10 @@ class Client(
             new_client._session_kwargs["message_handler"] = TaskNotificationHandler(
                 new_client
             )
+        # Rebind the task-status notification binding so it routes to the clone.
+        new_client._session_kwargs["notification_bindings"] = [
+            new_client._task_status_binding()
+        ]
 
         new_client.name += f":{secrets.token_hex(2)}"
 
@@ -778,8 +791,11 @@ class Client(
         Called when notifications/tasks/status is received from server.
         Updates Task object's cache and triggers events/callbacks.
         """
-        # Extract task ID from notification params
-        task_id = notification.params.task_id
+        self._handle_task_status_params(notification.params)
+
+    def _handle_task_status_params(self, params: TaskStatusNotificationParams) -> None:
+        """Route task status notification params to the matching Task object."""
+        task_id = params.task_id
         if not task_id:
             return
 
@@ -789,9 +805,29 @@ class Client(
             task = task_ref()  # Dereference weakref
             if task:
                 # Convert notification params to GetTaskResult (they share the same fields via Task)
-                status = GetTaskResult.model_validate(notification.params.model_dump())
+                status = GetTaskResult.model_validate(params.model_dump())
                 task._handle_status_notification(status)
 
+    def _task_status_binding(self) -> NotificationBinding[TaskStatusNotificationParams]:
+        """Build a binding routing `notifications/tasks/status` to Task objects.
+
+        SDK v2 drops notifications whose method is absent from the negotiated
+        version's core tables before they reach the message_handler; a binding is
+        the supported channel for observing such vendor notifications.
+        """
+        client_ref = weakref.ref(self)
+
+        async def _handler(params: TaskStatusNotificationParams) -> None:
+            client = client_ref()
+            if client is not None:
+                client._handle_task_status_params(params)
+
+        return NotificationBinding(
+            method="notifications/tasks/status",
+            params_type=TaskStatusNotificationParams,
+            handler=_handler,
+        )
+
     async def close(self):
         await self._disconnect(force=True)
         await self.transport.close()
diff --git a/fastmcp_slim/fastmcp/client/transports/base.py b/fastmcp_slim/fastmcp/client/transports/base.py
index bc13782be..87326413c 100644
--- a/fastmcp_slim/fastmcp/client/transports/base.py
+++ b/fastmcp_slim/fastmcp/client/transports/base.py
@@ -1,11 +1,12 @@
 import abc
 import contextlib
-from collections.abc import AsyncIterator
-from typing import Literal, TypeVar
+from collections.abc import AsyncIterator, Sequence
+from typing import Any, Literal, TypeVar
 
 import httpx
 import mcp_types
 from mcp import ClientSession
+from mcp.client.extension import NotificationBinding
 from mcp.client.session import (
     ElicitationFnT,
     ListRootsFnT,
@@ -30,6 +31,7 @@ class SessionKwargs(TypedDict, total=False):
     elicitation_callback: ElicitationFnT | None
     message_handler: MessageHandlerFnT | None
     client_info: mcp_types.Implementation | None
+    notification_bindings: Sequence[NotificationBinding[Any]] | None
 
 
 class ClientTransport(abc.ABC):