Deliver task status notifications: persist subscription task, nest params _meta (SDK v2)

fastmcp_slim/fastmcp/server/tasks/handlers.py

@@ -5,6 +5,7 @@ Handles queuing tool/prompt/resource executions to Docket as background tasks.
 
 from __future__ import annotations
 
+import asyncio
 import uuid
 from contextlib import suppress
 from datetime import datetime, timezone
@@ -180,22 +181,33 @@ async def submit_to_docket(
     else:
         await component.add_to_docket(docket, arguments, fn_key=key, task_key=task_key)  # type: ignore[call-arg]  # ty:ignore[invalid-argument-type, too-many-positional-arguments]
 
-    # Spawn subscription task to send status notifications (SEP-1686 optional feature)
-    # Start subscription in session's task group (persists for connection lifetime)
+    # Spawn subscription task to send status notifications (SEP-1686 optional feature).
+    # SDK v2 constructs a ServerSession per request and exposes no per-connection
+    # task group, so the subscription runs as a standalone asyncio task that
+    # outlives the submitting request; it is cancelled when the connection closes.
     # Deferred: subscriptions and notifications depend on docket at import time
     from fastmcp.server.tasks.subscriptions import subscribe_to_task_updates
 
-    if hasattr(ctx.session, "_subscription_task_group"):
-        tg = ctx.session._subscription_task_group
-        if tg:
-            tg.start_soon(  # type: ignore[union-attr]  # ty:ignore[unresolved-attribute]
-                subscribe_to_task_updates,
-                server_task_id,
-                task_key,
-                ctx.session,
-                docket,
-                poll_interval_ms,
-            )
+    subscription_task = asyncio.create_task(
+        subscribe_to_task_updates(
+            server_task_id,
+            task_key,
+            ctx.session,
+            docket,
+            poll_interval_ms,
+        ),
+        name=f"task-subscription-{server_task_id[:8]}",
+    )
+    connection = getattr(ctx.session, "_connection", None)
+    if connection is not None:
+
+        async def _cancel_subscription() -> None:
+            if not subscription_task.done():
+                subscription_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await subscription_task
+
+        connection.exit_stack.push_async_callback(_cancel_subscription)
 
     # Deferred: notifications depends on docket at import time
     from fastmcp.server.tasks.notifications import (

fastmcp_slim/fastmcp/server/tasks/notifications.py

@@ -186,11 +186,16 @@ async def _send_mcp_notification(
     if method != "notifications/tasks/status":
         raise ValueError(f"Unsupported notification method for subscriber: {method}")
 
+    # SDK v2: a notification's `_meta` lives on its params (`params._meta`), not
+    # at the notification envelope level, so nest it under params before parsing.
+    params_dict = dict(notification_dict.get("params", {}))
+    meta_dict = notification_dict.get("_meta")
+    if meta_dict is not None:
+        params_dict["_meta"] = meta_dict
     notification = mcp_types.TaskStatusNotification.model_validate(
         {
             "method": "notifications/tasks/status",
-            "params": notification_dict.get("params", {}),
-            "_meta": notification_dict.get("_meta"),
+            "params": params_dict,
         }
     )
     # SDK v2: `ServerNotification` is a union type; send the bare model.

tests/client/tasks/test_client_task_notifications.py

@@ -215,8 +215,9 @@ async def test_wait_returns_on_input_required(task_notification_server):
     async with Client(task_notification_server) as client:
         task = await client.call_tool("quick_task", {"value": 1}, task=True)
 
-        # Directly inject an input_required status into the cache and signal the event
-        now = datetime.now(timezone.utc)
+        # Directly inject an input_required status into the cache and signal the event.
+        # SDK v2 types the Task timestamps as ISO 8601 strings.
+        now = datetime.now(timezone.utc).isoformat()
         input_required_status = GetTaskResult(
             task_id=task._task_id,
             status="input_required",

tests/server/tasks/test_notifications.py

@@ -12,7 +12,6 @@ import mcp_types
 from fastmcp import FastMCP
 from fastmcp.client import Client
 from fastmcp.client.elicitation import ElicitResult
-from fastmcp.client.messages import MessageHandler
 from fastmcp.server.context import Context
 from fastmcp.server.elicitation import AcceptedElicitation
 from fastmcp.server.tasks.notifications import (
@@ -20,24 +19,6 @@ from fastmcp.server.tasks.notifications import (
 )
 
 
-class NotificationCaptureHandler(MessageHandler):
-    """Capture server notifications for test assertions."""
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.notifications: list[mcp_types.ServerNotification] = []
-
-    async def on_notification(self, message: mcp_types.ServerNotification) -> None:
-        self.notifications.append(message)
-
-    def for_method(self, method: str) -> list[mcp_types.ServerNotification]:
-        return [
-            notification
-            for notification in self.notifications
-            if notification.root.method == method
-        ]
-
-
 class TestNotificationIntegration:
     """Integration tests for the notification queue using real Docket memory backend.
 
@@ -51,12 +32,14 @@ class TestNotificationIntegration:
     async def test_notification_delivered_during_elicitation(self):
         """Full E2E: notification queue delivers input_required metadata to client.
 
-        The elicitation relay handles the response via the client's
-        elicitation_handler. We verify both the notification metadata
-        structure and the end-to-end elicitation flow.
+        SDK v2 does not carry `notifications/tasks/status` in any protocol
+        version's core notification tables, so it is delivered through the
+        client's task-status notification binding (routed to Task objects) rather
+        than the message_handler. We observe it via `on_status_change`, whose
+        GetTaskResult carries the notification's `_meta`.
         """
         mcp = FastMCP("notification-test")
-        notification_handler = NotificationCaptureHandler()
+        captured: list[mcp_types.GetTaskResult] = []
 
         @mcp.tool(task=True)
         async def elicit_tool(ctx: Context) -> str:
@@ -70,20 +53,19 @@ class TestNotificationIntegration:
 
         async with Client(
             mcp,
-            message_handler=notification_handler,
             elicitation_handler=elicitation_handler,
         ) as client:
             task = await client.call_tool("elicit_tool", {}, task=True)
+            task.on_status_change(captured.append)
 
             await task.wait(timeout=10.0)
             result = await task.result()
             assert result.data == "got: hello"
 
             # Verify the input_required notification was delivered with metadata
-            notification: mcp_types.ServerNotification | None = None
-            candidates = notification_handler.for_method("notifications/tasks/status")
-            for candidate in reversed(candidates):
-                candidate_meta = getattr(candidate.root, "_meta", None)
+            notification: mcp_types.GetTaskResult | None = None
+            for candidate in reversed(captured):
+                candidate_meta = candidate.meta
                 related_task = (
                     candidate_meta.get("io.modelcontextprotocol/related-task")
                     if isinstance(candidate_meta, dict)
@@ -97,7 +79,7 @@ class TestNotificationIntegration:
                     break
 
             assert notification is not None, "expected notifications/tasks/status"
-            task_meta = getattr(notification.root, "_meta", None)
+            task_meta = notification.meta
             assert isinstance(task_meta, dict)
 
             related_task = task_meta.get("io.modelcontextprotocol/related-task")