diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py index b778a63b2..0cdc9636a 100644 --- a/tests/server/tasks/test_context_background_task.py +++ b/tests/server/tasks/test_context_background_task.py @@ -1,11 +1,23 @@ -"""Tests for Context background task support (SEP-1686).""" +"""Tests for Context background task support (SEP-1686). + +Tests Context API surface (unit) and background task elicitation (integration). +Integration tests use Client(mcp) with the real memory:// Docket backend — +no mocking of Redis, Docket, or session internals. +""" + +import asyncio import pytest from fastmcp import FastMCP +from fastmcp.client import Client from fastmcp.server.context import Context -from fastmcp.server.elicitation import AcceptedElicitation -from fastmcp.server.tasks.elicitation import elicit_for_task, handle_task_input +from fastmcp.server.elicitation import AcceptedElicitation, DeclinedElicitation +from fastmcp.server.tasks.elicitation import handle_task_input + +# ============================================================================= +# Unit tests: Context API surface (no Redis/Docket needed) +# ============================================================================= class TestContextBackgroundTaskSupport: @@ -54,7 +66,6 @@ class TestContextSessionProperty: mock_session = MockSession() ctx = Context(mcp, session=mock_session, task_id="test-task-123") # type: ignore[arg-type] - # In background task mode, should return the stored session assert ctx.session is mock_session def test_session_uses_stored_session_during_on_initialize(self): @@ -65,23 +76,19 @@ class TestContextSessionProperty: _fastmcp_state_prefix = "test-session" mock_session = MockSession() - # Simulating on_initialize: has session but not a background task ctx = Context(mcp, session=mock_session) # type: ignore[arg-type] - # Should return the stored session as fallback assert ctx.session is mock_session class TestContextElicitBackgroundTask: """Tests for Context.elicit() in background task mode.""" - @pytest.mark.asyncio async def test_elicit_raises_when_background_task_but_no_docket(self): """elicit() should raise when in background task mode but Docket unavailable.""" mcp = FastMCP("test") ctx = Context(mcp, task_id="test-task-123") - # Set up minimal session mock class MockSession: _fastmcp_state_prefix = "test-session" @@ -91,6 +98,53 @@ class TestContextElicitBackgroundTask: await ctx.elicit("Need input", str) +class TestElicitFailFast: + """Tests for elicit_for_task fail-fast on notification push failure.""" + + async def test_elicit_returns_cancel_when_notification_push_fails(self): + """elicit_for_task should return cancel immediately when push_notification fails. + + If the client can't receive the input_required notification, waiting + for a response that will never come would block for up to 1 hour. + Instead, we return cancel immediately (fail-fast). + + This test patches ONLY push_notification — all other components + (Docket, Redis, session) are real via the memory:// backend. + """ + from unittest.mock import patch + + from fastmcp.server.elicitation import CancelledElicitation + + mcp = FastMCP("failfast-test") + elicit_started = asyncio.Event() + captured: dict[str, object] = {} + + @mcp.tool(task=True) + async def failfast_tool(ctx: Context) -> str: + elicit_started.set() + result = await ctx.elicit("This notification will fail", str) + captured["result_type"] = type(result).__name__ + captured["is_cancelled"] = isinstance(result, CancelledElicitation) + return "done" + + # Patch push_notification BEFORE starting client so it's active + # when the tool runs in the Docket worker + with patch( + "fastmcp.server.tasks.notifications.push_notification", + side_effect=ConnectionError("Redis queue unavailable"), + ): + async with Client(mcp) as client: + task = await client.call_tool("failfast_tool", {}, task=True) + await asyncio.wait_for(elicit_started.wait(), timeout=5.0) + await task.wait(timeout=10.0) + result = await task.result() + assert result.data == "done" + + # The tool should have received CancelledElicitation (fail-fast) + assert captured["is_cancelled"] is True + assert captured["result_type"] == "CancelledElicitation" + + class TestContextDocumentation: """Tests to verify Context documentation and API surface.""" @@ -110,735 +164,97 @@ class TestContextDocumentation: assert "background task" in Context.session.fget.__doc__.lower() -class TestBackgroundTaskElicitationE2E: - """End-to-end tests for background task elicitation (SEP-1686). +# ============================================================================= +# Integration tests: Client(mcp) + memory:// Docket backend +# ============================================================================= - These tests demonstrate the full flow: - 1. Client calls a tool with task=True (background execution) - 2. Tool uses ctx.elicit() to request user input - 3. Task status changes to "input_required" - 4. Client sends input via handle_task_input() - 5. Task resumes and completes with the elicited value - This simulates what a client would see when interacting with - a background task that needs user input. +class TestBackgroundTaskIntegration: + """Integration tests for background task context using real Docket memory backend. + + These tests use Client(mcp) with the memory:// broker — no mocking. + The memory:// backend provides a fully functional in-memory Redis store + that Docket uses automatically when running tests. """ - async def test_elicit_for_task_stores_request_in_redis(self): - """Test that elicit_for_task stores the elicitation request in Redis. - - This tests the Redis coordination layer that enables client interaction. - When a background task calls elicit(), the request is stored in Redis - so clients can retrieve it and respond. - """ - from unittest.mock import AsyncMock, MagicMock, patch - - from fastmcp.server.tasks.elicitation import ( - elicit_for_task, - ) - - # Create mocks - mock_redis = AsyncMock() - mock_redis.set = AsyncMock() - mock_redis.get = AsyncMock(return_value=None) # No response yet - mock_redis.delete = AsyncMock() - - mock_docket = MagicMock() - mock_docket.redis = MagicMock(return_value=AsyncMock()) - mock_docket.redis.return_value.__aenter__ = AsyncMock(return_value=mock_redis) - mock_docket.redis.return_value.__aexit__ = AsyncMock() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - mock_session = MagicMock() - mock_session._fastmcp_state_prefix = "test-session-id" - mock_session.send_notification = AsyncMock() - - # Call elicit_for_task with a short timeout to avoid blocking - with patch("fastmcp.server.tasks.elicitation.ELICIT_TTL_SECONDS", 1): - with patch("fastmcp.server.tasks.elicitation.asyncio.sleep", AsyncMock()): - # Make it return after first poll - mock_redis.get = AsyncMock( - return_value=b'{"action": "accept", "content": {"value": 42}}' - ) - - result = await elicit_for_task( - task_id="test-task-123", - session=mock_session, - message="Please provide a number", - schema={ - "type": "object", - "properties": {"value": {"type": "integer"}}, - }, - fastmcp=mock_fastmcp, - ) - - # Verify the result - assert result.action == "accept" - assert result.content == {"value": 42} - - # Verify Redis operations were called - assert mock_redis.set.call_count >= 2 # request + status - - async def test_handle_task_input_stores_response(self): - """Test that handle_task_input stores the response in Redis. - - This tests the client-side flow: when a client sends input via - tasks/sendInput, the response is stored in Redis for the waiting task. - """ - from unittest.mock import AsyncMock, MagicMock - - # Create mocks - mock_redis = AsyncMock() - mock_redis.get = AsyncMock(return_value=b"waiting") # Status is waiting - mock_redis.set = AsyncMock() - - mock_docket = MagicMock() - mock_docket.redis = MagicMock(return_value=AsyncMock()) - mock_docket.redis.return_value.__aenter__ = AsyncMock(return_value=mock_redis) - mock_docket.redis.return_value.__aexit__ = AsyncMock() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - # Call handle_task_input - success = await handle_task_input( - task_id="test-task-123", - session_id="test-session-id", - action="accept", - content={"value": 42}, - fastmcp=mock_fastmcp, - ) - - # Verify success - assert success is True - - # Verify Redis operations - assert mock_redis.set.call_count == 2 # response + status update - - async def test_handle_task_input_rejects_when_not_waiting(self): - """Test that handle_task_input rejects input when task isn't waiting. - - This verifies proper state management - clients can only send input - when a task is actually waiting for it. - """ - from unittest.mock import AsyncMock, MagicMock - - mock_redis = AsyncMock() - mock_redis.get = AsyncMock(return_value=None) # No waiting status - - mock_docket = MagicMock() - mock_docket.redis = MagicMock(return_value=AsyncMock()) - mock_docket.redis.return_value.__aenter__ = AsyncMock(return_value=mock_redis) - mock_docket.redis.return_value.__aexit__ = AsyncMock() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - success = await handle_task_input( - task_id="test-task-123", - session_id="test-session-id", - action="accept", - content={"value": 42}, - fastmcp=mock_fastmcp, - ) - - # Should fail because no task is waiting - assert success is False - - async def test_elicit_for_task_sends_notification(self): - """Test that elicit_for_task sends input_required notification. - - Per SEP-1686, the server should send notifications/tasks/updated - with status="input_required" when a task needs input. - """ - from unittest.mock import AsyncMock, MagicMock, patch - - mock_redis = AsyncMock() - mock_redis.set = AsyncMock() - mock_redis.get = AsyncMock( - return_value=b'{"action": "accept", "content": {"value": 1}}' - ) - mock_redis.delete = AsyncMock() - - mock_docket = MagicMock() - mock_docket.redis = MagicMock(return_value=AsyncMock()) - mock_docket.redis.return_value.__aenter__ = AsyncMock(return_value=mock_redis) - mock_docket.redis.return_value.__aexit__ = AsyncMock() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - mock_session = MagicMock() - mock_session._fastmcp_state_prefix = "test-session" - mock_session.send_notification = AsyncMock() - - with patch("fastmcp.server.tasks.elicitation.asyncio.sleep", AsyncMock()): - await elicit_for_task( - task_id="my-task-id", - session=mock_session, - message="Enter value", - schema={"type": "object"}, - fastmcp=mock_fastmcp, - ) - - # Verify notification was sent - mock_session.send_notification.assert_called_once() - notification = mock_session.send_notification.call_args[0][0] - assert notification.method == "notifications/tasks/updated" - - async def test_elicit_for_task_timeout_returns_cancel(self): - """Test that elicit_for_task returns cancel on timeout. - - If no response is received within the TTL, the elicitation - should be treated as cancelled. - """ - from unittest.mock import AsyncMock, MagicMock, patch - - mock_redis = AsyncMock() - mock_redis.set = AsyncMock() - mock_redis.get = AsyncMock(return_value=None) # Never responds - mock_redis.delete = AsyncMock() - - mock_docket = MagicMock() - mock_docket.redis = MagicMock(return_value=AsyncMock()) - mock_docket.redis.return_value.__aenter__ = AsyncMock(return_value=mock_redis) - mock_docket.redis.return_value.__aexit__ = AsyncMock() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - mock_session = MagicMock() - mock_session._fastmcp_state_prefix = "test-session" - mock_session.send_notification = AsyncMock() - - # Use very short TTL for test - with patch("fastmcp.server.tasks.elicitation.ELICIT_TTL_SECONDS", 0.1): - with patch( - "fastmcp.server.tasks.elicitation.asyncio.sleep", - AsyncMock(), - ): - result = await elicit_for_task( - task_id="timeout-task", - session=mock_session, - message="This will timeout", - schema={"type": "object"}, - fastmcp=mock_fastmcp, - ) - - # Should return cancel on timeout - assert result.action == "cancel" - assert result.content is None - - async def test_elicit_notification_includes_full_schema(self): - """Test that the notification includes the full JSON schema for complex types. - - This test demonstrates what the client sees when eliciting a Pydantic model. - The client receives a full JSON Schema that describes the expected input, - which they can use to: - - Render a dynamic form - - Validate user input before sending - - Show field descriptions to the user - - Example notification metadata for a UserInfo model: - ```json - { - "modelcontextprotocol.io/related-task": { - "taskId": "test-task", - "status": "input_required", - "statusMessage": "Please provide user info", - "elicitation": { - "requestId": "...", - "message": "Please provide user info", - "requestedSchema": { - "type": "object", - "properties": { - "name": {"type": "string", "title": "Name"}, - "age": {"type": "integer", "title": "Age"} - }, - "required": ["name", "age"], - "title": "UserInfo" - } - } - } - } - ``` - """ - from unittest.mock import AsyncMock, MagicMock, patch - - from pydantic import BaseModel - - class UserInfo(BaseModel): - """User information for registration.""" - - name: str - age: int - - mock_redis = AsyncMock() - mock_redis.set = AsyncMock() - mock_redis.get = AsyncMock( - return_value=b'{"action": "accept", "content": {"name": "Alice", "age": 30}}' - ) - mock_redis.delete = AsyncMock() - - mock_docket = MagicMock() - mock_docket.redis = MagicMock(return_value=AsyncMock()) - mock_docket.redis.return_value.__aenter__ = AsyncMock(return_value=mock_redis) - mock_docket.redis.return_value.__aexit__ = AsyncMock() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - mock_session = MagicMock() - mock_session._fastmcp_state_prefix = "test-session" - mock_session.send_notification = AsyncMock() - - # Create task-aware context - ctx = Context( - mock_fastmcp, - session=mock_session, - task_id="schema-test-task", - ) - - # Call elicit with a Pydantic model type - with patch("fastmcp.server.tasks.elicitation.asyncio.sleep", AsyncMock()): - result = await ctx.elicit("Please provide user info", UserInfo) - - # Verify the notification includes the full schema - mock_session.send_notification.assert_called_once() - notification = mock_session.send_notification.call_args[0][0] - meta = notification._meta - related_task = meta["modelcontextprotocol.io/related-task"] - schema = related_task["elicitation"]["requestedSchema"] - - # Verify schema structure matches UserInfo - assert schema["type"] == "object" - assert "properties" in schema - assert "name" in schema["properties"] - assert "age" in schema["properties"] - assert schema["properties"]["name"]["type"] == "string" - assert schema["properties"]["age"]["type"] == "integer" - assert "required" in schema - assert set(schema["required"]) == {"name", "age"} - - # Verify the result is properly parsed into the Pydantic model - assert result.action == "accept" - assert isinstance(result, AcceptedElicitation) # Type narrowing - assert isinstance(result.data, UserInfo) - assert result.data.name == "Alice" - assert result.data.age == 30 - - -class TestBackgroundTaskContextWiring: - """Integration tests for Context wiring in Docket workers. - - These tests verify that when a background task runs in a Docket worker, - the Context dependency is properly created with task_id and session, - allowing ctx.elicit() to work transparently. - - Per Chris Guidry's review request: "Could we get at least one test showing - the end-to-end of it working, with a background task that's eliciting input? - This will help with what the client-side sees when this happens." - - The key test is `test_context_elicit_full_flow_with_mocked_redis` which shows: - - CLIENT RECEIVES: - notifications/tasks/updated with: - - taskId: the background task ID - - status: "input_required" - - statusMessage: the elicit prompt - - elicitation.requestedSchema: JSON schema for expected input - - CLIENT RESPONDS: - handle_task_input(task_id, session_id, action="accept", content={...}) - - TOOL RECEIVES: - AcceptedElicitation(action="accept", data=) - """ - - async def test_context_is_created_with_task_id_in_worker(self): - """Test that Context is created with task_id when running in Docket worker. - - This verifies the wiring from _CurrentContext that creates a task-aware - Context when get_task_context() returns TaskContextInfo. - """ - from unittest.mock import MagicMock, patch - - from fastmcp.server.dependencies import ( - TaskContextInfo, - _current_server, - _CurrentContext, - _task_sessions, - ) - - # Set up mock server - mock_server = MagicMock() - mock_server._docket = MagicMock() - server_token = _current_server.set(MagicMock(return_value=mock_server)) - - # Set up mock session in registry - mock_session = MagicMock() - mock_session._fastmcp_state_prefix = "test-session-id" - _task_sessions["test-session-id"] = MagicMock(return_value=mock_session) - - try: - # Mock get_task_context to return TaskContextInfo - task_info = TaskContextInfo( - task_id="test-task-123", - session_id="test-session-id", - ) - with patch( - "fastmcp.server.dependencies.get_task_context", - return_value=task_info, - ): - # Create the dependency and enter it - dep = _CurrentContext() - ctx = await dep.__aenter__() - - # Verify context is task-aware - assert ctx.is_background_task is True - assert ctx.task_id == "test-task-123" - assert ctx.session is mock_session - - # Clean up - await dep.__aexit__(None, None, None) - finally: - _current_server.reset(server_token) - _task_sessions.pop("test-session-id", None) - - async def test_context_falls_back_to_foreground_mode(self): - """Test that Context uses foreground mode when not in worker context. - - When _current_context has a value (normal request handling), - _CurrentContext should return that context instead of creating a new one. - """ - from unittest.mock import MagicMock - - from fastmcp.server.context import Context, _current_context - from fastmcp.server.dependencies import _CurrentContext - - mcp = MagicMock() - foreground_ctx = Context(mcp) - - # Set the foreground context - token = _current_context.set(foreground_ctx) - try: - dep = _CurrentContext() - ctx = await dep.__aenter__() - - # Should return the foreground context - assert ctx is foreground_ctx - assert ctx.is_background_task is False - - await dep.__aexit__(None, None, None) - finally: - _current_context.reset(token) - - async def test_session_registered_when_task_submitted(self): - """Test that session is registered when a task is submitted to Docket. - - This verifies that submit_to_docket calls register_task_session, - which enables the Context wiring in background workers. - """ - import asyncio - - from fastmcp import FastMCP - from fastmcp.client import Client - from fastmcp.server.dependencies import get_task_session - - mcp = FastMCP("test-server") - - task_started = asyncio.Event() - session_id_captured = None + async def test_report_progress_in_background_task(self): + """report_progress() should complete without error in a background task.""" + mcp = FastMCP("progress-test") + progress_reported = asyncio.Event() @mcp.tool(task=True) - async def capture_session_tool(ctx: Context) -> str: - """Tool that captures the session ID for verification.""" - nonlocal session_id_captured - task_started.set() - # Access session to verify it works - session_id_captured = ctx.session_id + async def progress_tool(ctx: Context) -> str: + await ctx.report_progress(0, 100, "Starting...") + await ctx.report_progress(50, 100, "Half done") + await ctx.report_progress(100, 100, "Complete") + progress_reported.set() return "done" async with Client(mcp) as client: - # Start the task - task = await client.call_tool("capture_session_tool", {}, task=True) - assert task is not None - - # Wait for the task to start - await asyncio.wait_for(task_started.wait(), timeout=5.0) - - # Verify the session was registered - assert session_id_captured is not None - # The session should be retrievable via get_task_session - # (it was registered when the task was submitted) - # Session may be available or None if cleaned up - key is registration happened - _ = get_task_session(session_id_captured) - - # Wait for task to complete + task = await client.call_tool("progress_tool", {}, task=True) + await asyncio.wait_for(progress_reported.wait(), timeout=5.0) await task.wait(timeout=5.0) result = await task.result() assert result.data == "done" - async def test_context_elicit_works_in_background_task(self): - """E2E test: verify Context is properly wired in background tasks. - - This test demonstrates that: - 1. Context.task_id is set correctly in background tasks - 2. Context.is_background_task returns True - 3. Context.session_id is available - - The wiring is what enables ctx.elicit() to work in background tasks. - """ - import asyncio - - from fastmcp import FastMCP - from fastmcp.client import Client - from fastmcp.server.context import Context - - mcp = FastMCP("context-wiring-test") - - # Track what happens in the background task + async def test_context_wiring_in_background_task(self): + """Context should be properly wired with task_id and session_id.""" + mcp = FastMCP("wiring-test") task_completed = asyncio.Event() - captured_task_id: str | None = None - captured_session_id: str | None = None - captured_is_background: bool | None = None + captured: dict[str, object] = {} @mcp.tool(task=True) - async def verify_context_tool(ctx: Context) -> str: - """Tool that verifies Context is wired correctly for background tasks.""" - nonlocal captured_task_id, captured_session_id, captured_is_background - - # Capture context properties - this is the key verification - captured_task_id = ctx.task_id - captured_session_id = ctx.session_id - captured_is_background = ctx.is_background_task - + async def verify_wiring(ctx: Context) -> str: + captured["task_id"] = ctx.task_id + captured["session_id"] = ctx.session_id + captured["is_background"] = ctx.is_background_task task_completed.set() - return f"task_id={ctx.task_id}, is_background={ctx.is_background_task}" + return "ok" async with Client(mcp) as client: - # Start the background task - task = await client.call_tool("verify_context_tool", {}, task=True) - assert task is not None - assert task.task_id is not None - - # Wait for the task to complete - await asyncio.wait_for(task_completed.wait(), timeout=10.0) - - # Verify Context was properly wired in the background task - assert captured_task_id is not None, "Context.task_id should be set" - assert captured_session_id is not None, "Context.session_id should be set" - assert captured_is_background is True, ( - "Context.is_background_task should be True" - ) - - # Wait for task result - await task.wait(timeout=10.0) + task = await client.call_tool("verify_wiring", {}, task=True) + await asyncio.wait_for(task_completed.wait(), timeout=5.0) + await task.wait(timeout=5.0) result = await task.result() - assert "is_background=True" in result.data + assert result.data == "ok" - async def test_context_elicit_full_flow_with_mocked_redis(self): - """E2E test with mocked Redis to show complete elicitation flow. + assert captured["task_id"] is not None + assert captured["session_id"] is not None + assert captured["is_background"] is True - This test demonstrates what the client sees during background task - elicitation, with a mocked Redis layer to avoid requiring real Redis. + async def test_elicit_accept_flow(self): + """E2E: tool elicits input, client accepts, tool receives value. Flow: - 1. Tool calls ctx.elicit() in background task - 2. Elicitation stores request in Redis, sends input_required notification - 3. Simulated client sends response via handle_task_input() - 4. Tool receives response and completes - - This is the key test that fulfills Chris Guidry's request for an - "end-to-end test showing a background task that's eliciting input" - and demonstrates "what the client-side sees when this happens." + 1. Tool calls ctx.elicit("name?", str) — blocks waiting for input + 2. Client polls handle_task_input(action="accept", content={"value":"Bob"}) + 3. Tool resumes with AcceptedElicitation(data="Bob") """ - import asyncio - from unittest.mock import AsyncMock, MagicMock - - from fastmcp.server.context import Context - from fastmcp.server.tasks.elicitation import handle_task_input - - # Shared Redis storage that both elicit and handle_task_input will use - redis_storage: dict[str, bytes] = {} - - # Create a mock Redis that uses our shared storage - class MockRedis: - async def set( - self, key: str, value: str | bytes, ex: int | None = None - ) -> None: - redis_storage[key] = value.encode() if isinstance(value, str) else value - - async def get(self, key: str) -> bytes | None: - return redis_storage.get(key) - - async def delete(self, *keys: str) -> None: - for key in keys: - redis_storage.pop(key, None) - - mock_redis = MockRedis() - - # Create mock context manager for redis() - class MockRedisContext: - async def __aenter__(self): - return mock_redis - - async def __aexit__(self, *args): - pass - - mock_docket = MagicMock() - mock_docket.redis = lambda: MockRedisContext() - mock_docket.key = lambda k: k - - mock_fastmcp = MagicMock() - mock_fastmcp._docket = mock_docket - - mock_session = MagicMock() - mock_session._fastmcp_state_prefix = "test-session-123" - mock_session.send_notification = AsyncMock() - - # Create task-aware context (as would be created in background worker) - ctx = Context( - mock_fastmcp, - session=mock_session, - task_id="test-task-456", - ) - - # Verify context is properly configured for background task - assert ctx.is_background_task is True - assert ctx.task_id == "test-task-456" - - # Start elicit in a background task (simulating the Docket worker) - async def run_elicit(): - return await ctx.elicit("What is your name?", str) - - elicit_task = asyncio.create_task(run_elicit()) - - # Wait for elicit to store request and start polling - # The elicit_for_task function stores the request and sends notification - await asyncio.sleep(0.2) - - # ═══════════════════════════════════════════════════════════════════════ - # CLIENT PERSPECTIVE: What does the client see? - # ═══════════════════════════════════════════════════════════════════════ - - # 1. CLIENT RECEIVES: notifications/tasks/updated notification - mock_session.send_notification.assert_called() - notification = mock_session.send_notification.call_args[0][0] - assert notification.method == "notifications/tasks/updated" - - # 2. CLIENT INSPECTS: The notification metadata tells the client: - # - Which task needs input (taskId) - # - What status the task is in (input_required) - # - What message to display (statusMessage) - # - The schema for the expected response (elicitation.requestedSchema) - meta = notification._meta - related_task = meta["modelcontextprotocol.io/related-task"] - - assert related_task["taskId"] == "test-task-456" - assert related_task["status"] == "input_required" - assert related_task["statusMessage"] == "What is your name?" - assert "elicitation" in related_task - assert related_task["elicitation"]["message"] == "What is your name?" - assert "requestedSchema" in related_task["elicitation"] - - # 3. CLIENT RESPONDS: Send input via handle_task_input - # This is what a real client would do when it receives input_required - success = await handle_task_input( - task_id="test-task-456", - session_id="test-session-123", - action="accept", - content={"value": "Alice"}, - fastmcp=mock_fastmcp, - ) - assert success is True, "Client should successfully send input" - - # ═══════════════════════════════════════════════════════════════════════ - # TOOL PERSPECTIVE: What does the tool receive? - # ═══════════════════════════════════════════════════════════════════════ - - # Wait for elicit to receive the response and return - result = await asyncio.wait_for(elicit_task, timeout=5.0) - - # Verify the result contains what the client sent - # AcceptedElicitation has 'action' and 'data' attributes - assert result.action == "accept" - assert result.data == "Alice" # The value from content["value"] - - async def test_context_elicit_with_real_docket_memory_backend(self): - """E2E test using Docket's real memory:// backend. - - This test uses the real Docket memory backend instead of mocking Redis, - as suggested by Chris Guidry during code review. The memory:// backend - provides a fully functional in-memory Redis-like store that Docket uses - automatically when running tests. - - Flow: - 1. Create FastMCP server with task-enabled tool that calls ctx.elicit() - 2. Start the task via Client (which initializes Docket with memory://) - 3. Background task blocks waiting for client input - 4. Simulate client sending input via handle_task_input() - 5. Task resumes and completes with the elicited value - - This demonstrates the complete elicitation flow with real infrastructure. - """ - import asyncio - - from fastmcp import FastMCP - from fastmcp.client import Client - from fastmcp.server.context import Context - from fastmcp.server.tasks.elicitation import handle_task_input - - mcp = FastMCP("elicit-memory-test") - - # Track task state using mutable container (avoids nonlocal) + mcp = FastMCP("elicit-accept-test") elicit_started = asyncio.Event() captured: dict[str, str | None] = {"task_id": None, "session_id": None} @mcp.tool(task=True) - async def ask_for_name(ctx: Context) -> str: - """Tool that elicits user's name via background task.""" - # Capture IDs for handle_task_input call + async def ask_name(ctx: Context) -> str: captured["task_id"] = ctx.task_id captured["session_id"] = ctx.session_id elicit_started.set() - # This will block until client sends input result = await ctx.elicit("What is your name?", str) - if isinstance(result, AcceptedElicitation): return f"Hello, {result.data}!" - else: - return "Elicitation was declined or cancelled" + return "No name provided" async with Client(mcp) as client: - # Start the background task - task = await client.call_tool("ask_for_name", {}, task=True) - assert task is not None - assert task.task_id is not None - - # Wait for task to reach elicit() call + task = await client.call_tool("ask_name", {}, task=True) await asyncio.wait_for(elicit_started.wait(), timeout=5.0) - # Poll until handle_task_input succeeds - # We need to wait for elicit_for_task to store the "waiting" status in Redis - # before we can send input. Using fixed-interval polling (not exponential - # backoff) because we're waiting for state, not recovering from errors. assert captured["task_id"] is not None assert captured["session_id"] is not None - max_attempts = 40 - poll_interval_seconds = 0.05 # 50ms - fast for tests, 2s max total + # Poll until the "waiting" status is stored in Redis success = False - for _ in range(max_attempts): + for _ in range(40): success = await handle_task_input( task_id=captured["task_id"], session_id=captured["session_id"], @@ -848,15 +264,127 @@ class TestBackgroundTaskContextWiring: ) if success: break - await asyncio.sleep(poll_interval_seconds) + await asyncio.sleep(0.05) - assert success is True, ( - f"handle_task_input should succeed within {max_attempts * poll_interval_seconds}s" - ) + assert success is True, "handle_task_input should succeed within 2s" - # Wait for task to complete await task.wait(timeout=10.0) result = await task.result() - - # Verify the tool received the elicited value and returned correctly assert result.data == "Hello, Bob!" + + async def test_elicit_decline_flow(self): + """E2E: tool elicits input, client declines, tool gets DeclinedElicitation.""" + mcp = FastMCP("elicit-decline-test") + elicit_started = asyncio.Event() + captured: dict[str, str | None] = {"task_id": None, "session_id": None} + + @mcp.tool(task=True) + async def optional_input(ctx: Context) -> str: + captured["task_id"] = ctx.task_id + captured["session_id"] = ctx.session_id + elicit_started.set() + + result = await ctx.elicit("Want to provide a name?", str) + if isinstance(result, DeclinedElicitation): + return "User declined" + if isinstance(result, AcceptedElicitation): + return f"Got: {result.data}" + return "Cancelled" + + async with Client(mcp) as client: + task = await client.call_tool("optional_input", {}, task=True) + await asyncio.wait_for(elicit_started.wait(), timeout=5.0) + + assert captured["task_id"] is not None + assert captured["session_id"] is not None + + success = False + for _ in range(40): + success = await handle_task_input( + task_id=captured["task_id"], + session_id=captured["session_id"], + action="decline", + content=None, + fastmcp=mcp, + ) + if success: + break + await asyncio.sleep(0.05) + + assert success is True + + await task.wait(timeout=10.0) + result = await task.result() + assert result.data == "User declined" + + async def test_elicit_with_pydantic_model(self): + """E2E: tool elicits structured Pydantic input, data round-trips correctly.""" + from pydantic import BaseModel + + class UserInfo(BaseModel): + name: str + age: int + + mcp = FastMCP("elicit-pydantic-test") + elicit_started = asyncio.Event() + captured: dict[str, str | None] = {"task_id": None, "session_id": None} + + @mcp.tool(task=True) + async def get_user_info(ctx: Context) -> str: + captured["task_id"] = ctx.task_id + captured["session_id"] = ctx.session_id + elicit_started.set() + + result = await ctx.elicit("Provide user info", UserInfo) + if isinstance(result, AcceptedElicitation): + assert isinstance(result.data, UserInfo) + return f"{result.data.name} is {result.data.age}" + return "No info" + + async with Client(mcp) as client: + task = await client.call_tool("get_user_info", {}, task=True) + await asyncio.wait_for(elicit_started.wait(), timeout=5.0) + + assert captured["task_id"] is not None + assert captured["session_id"] is not None + + success = False + for _ in range(40): + success = await handle_task_input( + task_id=captured["task_id"], + session_id=captured["session_id"], + action="accept", + content={"name": "Alice", "age": 30}, + fastmcp=mcp, + ) + if success: + break + await asyncio.sleep(0.05) + + assert success is True + + await task.wait(timeout=10.0) + result = await task.result() + assert result.data == "Alice is 30" + + async def test_handle_task_input_rejects_when_not_waiting(self): + """handle_task_input returns False when no task is waiting for input.""" + mcp = FastMCP("reject-test") + + @mcp.tool(task=True) + async def simple_tool() -> str: + return "done" + + async with Client(mcp) as client: + task = await client.call_tool("simple_tool", {}, task=True) + await task.wait(timeout=5.0) + + # Task already completed — no elicitation waiting + success = await handle_task_input( + task_id=task.task_id, + session_id="nonexistent-session", + action="accept", + content={"value": "too late"}, + fastmcp=mcp, + ) + assert success is False diff --git a/tests/server/tasks/test_notifications.py b/tests/server/tasks/test_notifications.py new file mode 100644 index 000000000..d669a5ed6 --- /dev/null +++ b/tests/server/tasks/test_notifications.py @@ -0,0 +1,104 @@ +"""Tests for distributed notification queue (SEP-1686). + +Integration tests verify that the notification queue works end-to-end +using Client(mcp) with the real memory:// Docket backend. +No mocking of Redis, sessions, or Docket internals. +""" + +import asyncio + +from fastmcp import FastMCP +from fastmcp.client import Client +from fastmcp.server.context import Context +from fastmcp.server.elicitation import AcceptedElicitation +from fastmcp.server.tasks.elicitation import handle_task_input +from fastmcp.server.tasks.notifications import ( + get_subscriber_count, +) + + +class TestNotificationIntegration: + """Integration tests for the notification queue using real Docket memory backend. + + The elicitation flow implicitly validates the full notification pipeline: + 1. Tool calls ctx.elicit() → stores request in Redis → pushes notification + 2. Subscriber picks up notification → sends MCP notification to client + 3. Client calls handle_task_input() → LPUSH response → BLPOP wakes tool + """ + + async def test_notification_delivered_during_elicitation(self): + """Full E2E: notification queue delivers elicitation notification to client.""" + mcp = FastMCP("notification-test") + elicit_started = asyncio.Event() + captured: dict[str, str | None] = {"task_id": None, "session_id": None} + + @mcp.tool(task=True) + async def elicit_tool(ctx: Context) -> str: + captured["task_id"] = ctx.task_id + captured["session_id"] = ctx.session_id + elicit_started.set() + + result = await ctx.elicit("Enter value", str) + if isinstance(result, AcceptedElicitation): + return f"got: {result.data}" + return "no value" + + async with Client(mcp) 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 + success = False + for _ in range(40): + success = await handle_task_input( + task_id=captured["task_id"], + session_id=captured["session_id"], + action="accept", + content={"value": "hello"}, + fastmcp=mcp, + ) + if success: + break + await asyncio.sleep(0.05) + + assert success is True + + await task.wait(timeout=10.0) + result = await task.result() + assert result.data == "got: hello" + + async def test_subscriber_lifecycle(self): + """Subscriber starts during background task and stops when client disconnects.""" + mcp = FastMCP("subscriber-test") + tool_started = asyncio.Event() + tool_continue = asyncio.Event() + + @mcp.tool(task=True) + async def lifecycle_tool(ctx: Context) -> str: + tool_started.set() + await asyncio.wait_for(tool_continue.wait(), timeout=10.0) + return "done" + + count_before = get_subscriber_count() + + async with Client(mcp) as client: + task = await client.call_tool("lifecycle_tool", {}, task=True) + await asyncio.wait_for(tool_started.wait(), timeout=5.0) + + # While a background task is running, subscriber should be active + count_during = get_subscriber_count() + assert count_during > count_before + + # Let the tool complete + tool_continue.set() + await task.wait(timeout=5.0) + result = await task.result() + assert result.data == "done" + + # After client disconnects, subscriber should be cleaned up + # Allow brief time for async cleanup + for _ in range(20): + if get_subscriber_count() == count_before: + break + await asyncio.sleep(0.05) + assert get_subscriber_count() == count_before