diff --git a/fastmcp_slim/fastmcp/server/tasks/handlers.py b/fastmcp_slim/fastmcp/server/tasks/handlers.py index d74f07f22..2a5577f4d 100644 --- a/fastmcp_slim/fastmcp/server/tasks/handlers.py +++ b/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 ( diff --git a/fastmcp_slim/fastmcp/server/tasks/notifications.py b/fastmcp_slim/fastmcp/server/tasks/notifications.py index 78a9034d4..9a662cd95 100644 --- a/fastmcp_slim/fastmcp/server/tasks/notifications.py +++ b/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. diff --git a/tests/client/tasks/test_client_task_notifications.py b/tests/client/tasks/test_client_task_notifications.py index b4debe115..94f7e1db3 100644 --- a/tests/client/tasks/test_client_task_notifications.py +++ b/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", diff --git a/tests/server/tasks/test_notifications.py b/tests/server/tasks/test_notifications.py index f6320b0d6..ca7abfe1a 100644 --- a/tests/server/tasks/test_notifications.py +++ b/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")