From 3d665d48c0763a06e080653c15324578fd46d2a3 Mon Sep 17 00:00:00 2001 From: Guillaume FORTAINE Date: Sun, 8 Feb 2026 15:20:22 +0100 Subject: [PATCH] fix: stabilize task notification integration tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🤖 Generated with Codex --- src/fastmcp/server/tasks/elicitation.py | 32 ++++++--- src/fastmcp/server/tasks/handlers.py | 32 ++++++--- src/fastmcp/server/tasks/notifications.py | 32 +++++---- .../tasks/test_context_background_task.py | 12 ++-- tests/server/tasks/test_notifications.py | 69 +++++++++++++++++-- tests/server/tasks/test_task_protocol.py | 2 +- 6 files changed, 134 insertions(+), 45 deletions(-) diff --git a/src/fastmcp/server/tasks/elicitation.py b/src/fastmcp/server/tasks/elicitation.py index b12d0f10e..299382df6 100644 --- a/src/fastmcp/server/tasks/elicitation.py +++ b/src/fastmcp/server/tasks/elicitation.py @@ -5,7 +5,7 @@ in Docket workers. Unlike regular MCP requests, background tasks don't have an active request context, so elicitation requires special handling: 1. Set task status to "input_required" via Redis -2. Send notifications/tasks/updated with elicitation metadata +2. Send notifications/tasks/status with elicitation metadata 3. Wait for client to send input via tasks/sendInput 4. Resume task execution with the provided input @@ -18,7 +18,8 @@ from __future__ import annotations import json import logging import uuid -from typing import TYPE_CHECKING, Any +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, cast import mcp.types from mcp import ServerSession @@ -115,15 +116,23 @@ async def elicit_for_task( ex=ELICIT_TTL_SECONDS, ) - # Send task status update notification with input_required status - # This follows SEP-1686 for background task status updates + # Send task status update notification with input_required status. + # Use notifications/tasks/status so typed MCP clients can consume it. # # NOTE: We use the distributed notification queue instead of session.send_notification() # This enables notifications to work when workers run in separate processes # (Azure Web PubSub / Service Bus inspired pattern) + timestamp = datetime.now(timezone.utc).isoformat() notification_dict = { - "method": "notifications/tasks/updated", - "params": {}, + "method": "notifications/tasks/status", + "params": { + "taskId": task_id, + "status": "input_required", + "statusMessage": message, + "createdAt": timestamp, + "lastUpdatedAt": timestamp, + "ttl": ELICIT_TTL_SECONDS * 1000, + }, "_meta": { "modelcontextprotocol.io/related-task": { "taskId": task_id, @@ -172,9 +181,12 @@ async def elicit_for_task( async with docket.redis() as redis: # BLPOP blocks until an item is pushed to the list or timeout # Returns tuple of (key, value) or None on timeout - result = await redis.blpop( - docket.key(response_key), - timeout=max_wait_seconds, + result = await cast( + Any, + redis.blpop( + [docket.key(response_key)], + timeout=max_wait_seconds, + ), ) if result: @@ -261,7 +273,7 @@ async def handle_task_input( # Push response to list - this wakes up the BLPOP in elicit_for_task # Using LPUSH instead of SET enables the efficient blocking wait pattern - await redis.lpush( + await redis.lpush( # type: ignore[invalid-await] # redis-py union type (sync/async) docket.key(response_key), json.dumps(response), ) diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index 646cf82dc..494fce87f 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/src/fastmcp/server/tasks/handlers.py @@ -112,21 +112,31 @@ async def submit_to_docket( register_task_session(session_id, ctx.session) - # Send notifications/tasks/created per SEP-1686 (mandatory) - # Send BEFORE queuing to avoid race where task completes before notification - notification = mcp.types.JSONRPCNotification( - jsonrpc="2.0", - method="notifications/tasks/created", - params={}, # Empty params per spec - _meta={ # type: ignore[call-arg] # _meta is Pydantic alias for meta field - "modelcontextprotocol.io/related-task": { + # Send an initial tasks/status notification before queueing. + # This guarantees clients can observe task creation immediately. + notification = mcp.types.TaskStatusNotification.model_validate( + { + "method": "notifications/tasks/status", + "params": { "taskId": server_task_id, - } - }, + "status": "working", + "statusMessage": "Task submitted", + "createdAt": created_at, + "lastUpdatedAt": created_at, + "ttl": ttl_ms, + "pollInterval": poll_interval_ms, + }, + "_meta": { + "modelcontextprotocol.io/related-task": { + "taskId": server_task_id, + } + }, + } ) + server_notification = mcp.types.ServerNotification(notification) with suppress(Exception): # Don't let notification failures break task creation - await ctx.session.send_notification(notification) # type: ignore[arg-type] + await ctx.session.send_notification(server_notification) # Queue function to Docket by key (result storage via execution_ttl) # Use component.add_to_docket() which handles calling conventions diff --git a/src/fastmcp/server/tasks/notifications.py b/src/fastmcp/server/tasks/notifications.py index 07fe95ff6..2c65e31e1 100644 --- a/src/fastmcp/server/tasks/notifications.py +++ b/src/fastmcp/server/tasks/notifications.py @@ -23,7 +23,7 @@ import logging import weakref from contextlib import suppress from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import mcp.types @@ -67,7 +67,7 @@ async def push_notification( } ) async with docket.redis() as redis: - await redis.lpush(key, message) + await redis.lpush(key, message) # type: ignore[invalid-await] # redis-py union type (sync/async) await redis.expire(key, NOTIFICATION_TTL_SECONDS) @@ -104,8 +104,8 @@ async def notification_subscriber_loop( # Blocking wait for notification (timeout refreshes heartbeat) # Using BRPOP (right pop) for FIFO order with LPUSH (left push) - result = await redis.brpop( - queue_key, timeout=SUBSCRIBER_TIMEOUT_SECONDS + result = await cast( + Any, redis.brpop([queue_key], timeout=SUBSCRIBER_TIMEOUT_SECONDS) ) if not result: continue # Timeout - refresh heartbeat and retry @@ -129,7 +129,7 @@ async def notification_subscriber_loop( # Re-queue with incremented attempt (back of queue) message["attempt"] = attempt + 1 message["last_error"] = str(send_error) - await redis.lpush(queue_key, json.dumps(message)) + await redis.lpush(queue_key, json.dumps(message)) # type: ignore[invalid-await] logger.debug( "Requeued notification for session %s (attempt %d): %s", session_id, @@ -166,18 +166,20 @@ async def _send_mcp_notification( session: MCP ServerSession notification_dict: Notification as dict (method, params, _meta) """ - # Build JSONRPCNotification from dict - notification = mcp.types.JSONRPCNotification( - jsonrpc="2.0", - method=notification_dict.get("method", "notifications/tasks/updated"), - params=notification_dict.get("params", {}), + method = notification_dict.get("method", "notifications/tasks/status") + if method != "notifications/tasks/status": + raise ValueError(f"Unsupported notification method for subscriber: {method}") + + notification = mcp.types.TaskStatusNotification.model_validate( + { + "method": "notifications/tasks/status", + "params": notification_dict.get("params", {}), + "_meta": notification_dict.get("_meta"), + } ) + server_notification = mcp.types.ServerNotification(notification) - # Preserve _meta if present (contains related-task info for elicitation) - if "_meta" in notification_dict: - notification._meta = notification_dict["_meta"] # type: ignore[attr-defined] - - await session.send_notification(notification) # type: ignore[arg-type] + await session.send_notification(server_notification) # ============================================================================= diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py index 0cdc9636a..219a7ae21 100644 --- a/tests/server/tasks/test_context_background_task.py +++ b/tests/server/tasks/test_context_background_task.py @@ -6,8 +6,10 @@ no mocking of Redis, Docket, or session internals. """ import asyncio +from typing import cast import pytest +from mcp import ServerSession from fastmcp import FastMCP from fastmcp.client import Client @@ -42,7 +44,7 @@ class TestContextBackgroundTaskSupport: mcp = FastMCP("test") ctx = Context(mcp, task_id="test-task-123") with pytest.raises(AttributeError): - ctx.task_id = "new-id" # type: ignore[misc] + setattr(ctx, "task_id", "new-id") class TestContextSessionProperty: @@ -64,7 +66,9 @@ class TestContextSessionProperty: _fastmcp_state_prefix = "test-session" mock_session = MockSession() - ctx = Context(mcp, session=mock_session, task_id="test-task-123") # type: ignore[arg-type] + ctx = Context( + mcp, session=cast(ServerSession, mock_session), task_id="test-task-123" + ) assert ctx.session is mock_session @@ -76,7 +80,7 @@ class TestContextSessionProperty: _fastmcp_state_prefix = "test-session" mock_session = MockSession() - ctx = Context(mcp, session=mock_session) # type: ignore[arg-type] + ctx = Context(mcp, session=cast(ServerSession, mock_session)) assert ctx.session is mock_session @@ -92,7 +96,7 @@ class TestContextElicitBackgroundTask: class MockSession: _fastmcp_state_prefix = "test-session" - ctx._session = MockSession() # type: ignore[assignment] + ctx._session = cast(ServerSession, MockSession()) with pytest.raises(RuntimeError, match="Docket"): await ctx.elicit("Need input", str) diff --git a/tests/server/tasks/test_notifications.py b/tests/server/tasks/test_notifications.py index d669a5ed6..07b2b7f3b 100644 --- a/tests/server/tasks/test_notifications.py +++ b/tests/server/tasks/test_notifications.py @@ -7,8 +7,11 @@ No mocking of Redis, sessions, or Docket internals. import asyncio +import mcp.types + from fastmcp import FastMCP from fastmcp.client import Client +from fastmcp.client.messages import MessageHandler from fastmcp.server.context import Context from fastmcp.server.elicitation import AcceptedElicitation from fastmcp.server.tasks.elicitation import handle_task_input @@ -17,6 +20,24 @@ 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. @@ -27,8 +48,9 @@ class TestNotificationIntegration: """ async def test_notification_delivered_during_elicitation(self): - """Full E2E: notification queue delivers elicitation notification to client.""" + """Full E2E: notification queue delivers input_required metadata to client.""" mcp = FastMCP("notification-test") + notification_handler = NotificationCaptureHandler() elicit_started = asyncio.Event() captured: dict[str, str | None] = {"task_id": None, "session_id": None} @@ -43,11 +65,50 @@ class TestNotificationIntegration: return f"got: {result.data}" return "no value" - async with Client(mcp) as client: + async with Client(mcp, message_handler=notification_handler) as client: task = await client.call_tool("elicit_tool", {}, task=True) await asyncio.wait_for(elicit_started.wait(), timeout=5.0) - # Poll until the task is waiting for input + assert captured["task_id"] is not None + assert captured["session_id"] is not None + + notification: mcp.types.ServerNotification | None = None + for _ in range(40): + candidates = notification_handler.for_method( + "notifications/tasks/status" + ) + for candidate in reversed(candidates): + candidate_meta = getattr(candidate.root, "_meta", None) + related_task = ( + candidate_meta.get("modelcontextprotocol.io/related-task") + if isinstance(candidate_meta, dict) + else None + ) + if ( + isinstance(related_task, dict) + and related_task.get("status") == "input_required" + ): + notification = candidate + break + if notification is not None: + break + await asyncio.sleep(0.05) + + assert notification is not None, "expected notifications/tasks/status" + task_meta = getattr(notification.root, "_meta", None) + assert isinstance(task_meta, dict) + + related_task = task_meta.get("modelcontextprotocol.io/related-task") + assert isinstance(related_task, dict) + assert related_task.get("taskId") == captured["task_id"] + assert related_task.get("status") == "input_required" + + elicitation = related_task.get("elicitation") + assert isinstance(elicitation, dict) + assert elicitation.get("message") == "Enter value" + assert isinstance(elicitation.get("requestId"), str) + assert isinstance(elicitation.get("requestedSchema"), dict) + success = False for _ in range(40): success = await handle_task_input( @@ -67,7 +128,7 @@ class TestNotificationIntegration: result = await task.result() assert result.data == "got: hello" - async def test_subscriber_lifecycle(self): + async def test_subscriber_started_and_cleaned_up(self): """Subscriber starts during background task and stops when client disconnects.""" mcp = FastMCP("subscriber-test") tool_started = asyncio.Event() diff --git a/tests/server/tasks/test_task_protocol.py b/tests/server/tasks/test_task_protocol.py index 1d1daa02c..9c6a88d22 100644 --- a/tests/server/tasks/test_task_protocol.py +++ b/tests/server/tasks/test_task_protocol.py @@ -48,7 +48,7 @@ async def test_task_metadata_includes_task_id_and_ttl(task_enabled_server): async def test_task_notification_sent_after_submission(task_enabled_server): - """Server sends notifications/tasks/created after task submission.""" + """Server sends an initial task status notification after submission.""" @task_enabled_server.tool(task=True) async def background_tool(message: str) -> str: