From 08974e50d94b6b856aefccfec30d4a1f82bd8715 Mon Sep 17 00:00:00 2001 From: Guillaume FORTAINE Date: Tue, 3 Feb 2026 01:28:42 +0100 Subject: [PATCH] feat(context): Add background task support for Context (SEP-1686) (#2905) --- src/fastmcp/server/context.py | 133 ++- src/fastmcp/server/dependencies.py | 140 ++- src/fastmcp/server/tasks/__init__.py | 3 + src/fastmcp/server/tasks/config.py | 18 +- src/fastmcp/server/tasks/elicitation.py | 229 +++++ src/fastmcp/server/tasks/handlers.py | 8 + .../tasks/test_context_background_task.py | 862 ++++++++++++++++++ 7 files changed, 1362 insertions(+), 31 deletions(-) create mode 100644 src/fastmcp/server/tasks/elicitation.py create mode 100644 tests/server/tasks/test_context_background_task.py diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 9a5eb3087..a9245fa2b 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -182,10 +182,45 @@ class Context: # Default TTL for session state: 1 day in seconds _STATE_TTL_SECONDS: int = 86400 - def __init__(self, fastmcp: FastMCP, session: ServerSession | None = None): + def __init__( + self, + fastmcp: FastMCP, + session: ServerSession | None = None, + *, + task_id: str | None = None, + ): self._fastmcp: weakref.ref[FastMCP] = weakref.ref(fastmcp) self._session: ServerSession | None = session # For state ops during init self._tokens: list[Token] = [] + # Background task support (SEP-1686) + self._task_id: str | None = task_id + + @property + def is_background_task(self) -> bool: + """True when this context is running in a background task (Docket worker). + + When True, certain operations like elicit() and sample() will use + task-aware implementations that can pause the task and wait for + client input. + + Example: + ```python + @server.tool(task=True) + async def my_task(ctx: Context) -> str: + # Works transparently in both foreground and background task modes + result = await ctx.elicit("Need input", str) + return str(result) + ``` + """ + return self._task_id is not None + + @property + def task_id(self) -> str | None: + """Get the background task ID if running in a background task. + + Returns None if not running in a background task context. + """ + return self._task_id @property def fastmcp(self) -> FastMCP: @@ -566,14 +601,27 @@ class Context: def session(self) -> ServerSession: """Access to the underlying session for advanced usage. - Raises RuntimeError if MCP request context is not available. + In request mode: Returns the session from the active request context. + In background task mode: Returns the session stored at Context creation. + + Raises RuntimeError if no session is available. """ - if self.request_context is None: - raise RuntimeError( - "session is not available because the MCP session has not been established yet. " - "Check `context.request_context` for None before accessing this attribute." - ) - return self.request_context.session + # Background task mode: use the stored session + if self.is_background_task and self._session is not None: + return self._session + + # Request mode: use request context + if self.request_context is not None: + return self.request_context.session + + # Fallback to stored session (e.g., during on_initialize) + if self._session is not None: + return self._session + + raise RuntimeError( + "session is not available because the MCP session has not been established yet. " + "Check `context.request_context` for None before accessing this attribute." + ) # Convenience methods for common log levels async def debug( @@ -841,7 +889,13 @@ class Context: - .text: The text representation (raw text or JSON for structured) - .result: The typed result (str for text, parsed object for structured) - .history: All messages exchanged during sampling + + Note: + Background task support for sampling is planned for a future release. + Currently, sampling in background tasks requires using the low-level + session.create_message() API directly. """ + # TODO: Add background task support similar to elicit() when is_background_task return await sample_impl( self, messages=messages, @@ -960,14 +1014,27 @@ class Context: response_type: The type of the response, which should be a primitive type or dataclass or BaseModel. If it is a primitive type, an object schema with a single "value" field will be generated. + + Note: + This method works transparently in both request and background task + contexts. In background task mode (SEP-1686), it will set the task + status to "input_required" and wait for the client to provide input. """ config = parse_elicit_response_type(response_type) - result = await self.session.elicit( - message=message, - requestedSchema=config.schema, - related_request_id=self.request_id, - ) + if self.is_background_task: + # Background task mode: use task-aware elicitation + result = await self._elicit_for_task( + message=message, + schema=config.schema, + ) + else: + # Standard request mode: use session.elicit directly + result = await self.session.elicit( + message=message, + requestedSchema=config.schema, + related_request_id=self.request_id, + ) if result.action == "accept": return handle_elicit_accept(config, result.content) @@ -978,6 +1045,46 @@ class Context: else: raise ValueError(f"Unexpected elicitation action: {result.action}") + async def _elicit_for_task( + self, + message: str, + schema: dict[str, Any], + ) -> mcp.types.ElicitResult: + """Send an elicitation request from a background task (SEP-1686). + + This method handles elicitation when running in a Docket worker context, + where there's no active MCP request. It: + 1. Sets the task status to "input_required" + 2. Sends the elicitation request with task metadata + 3. Waits for the client to provide input via tasks/sendInput + 4. Returns the result and resumes task execution + + Args: + message: The message to display to the user + schema: The JSON schema for the expected response + + Returns: + ElicitResult with the user's response + + Raises: + RuntimeError: If not running in a background task context + """ + if not self.is_background_task: + raise RuntimeError( + "_elicit_for_task called but not in a background task context" + ) + + # Import here to avoid circular imports and optional dependency issues + from fastmcp.server.tasks.elicitation import elicit_for_task + + return await elicit_for_task( + task_id=self._task_id, # type: ignore[arg-type] + session=self.session, + message=message, + schema=schema, + fastmcp=self.fastmcp, + ) + def _make_state_key(self, key: str) -> str: """Create session-prefixed key for state storage.""" return f"{self.session_id}:{key}" diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index 1f2032918..ce964fd31 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -13,6 +13,7 @@ import weakref from collections.abc import AsyncGenerator, Callable from contextlib import AsyncExitStack, asynccontextmanager from contextvars import ContextVar +from dataclasses import dataclass from functools import lru_cache from typing import TYPE_CHECKING, Any, Protocol, cast, get_type_hints, runtime_checkable @@ -35,6 +36,7 @@ from fastmcp.utilities.types import find_kwarg_by_type, is_class_member_of_type if TYPE_CHECKING: from docket import Docket from docket.worker import Worker + from mcp.server.session import ServerSession from fastmcp.server.context import Context from fastmcp.server.server import FastMCP @@ -50,12 +52,16 @@ __all__ = [ "CurrentRequest", "CurrentWorker", "Progress", + "TaskContextInfo", "get_access_token", "get_context", "get_http_headers", "get_http_request", "get_server", + "get_task_context", + "get_task_session", "is_docket_available", + "register_task_session", "require_docket", "resolve_dependencies", "transform_context_annotations", @@ -63,6 +69,95 @@ __all__ = [ ] +# --- TaskContextInfo and get_task_context --- + + +@dataclass(frozen=True, slots=True) +class TaskContextInfo: + """Information about the current background task context. + + Returned by ``get_task_context()`` when running inside a Docket worker. + Contains identifiers needed to communicate with the MCP session. + """ + + task_id: str + """The MCP task ID (server-generated UUID).""" + + session_id: str + """The session ID that submitted this task.""" + + +def get_task_context() -> TaskContextInfo | None: + """Get the current task context if running inside a background task worker. + + This function extracts task information from the Docket execution context. + Returns None if not running in a task context (e.g., foreground execution). + + Returns: + TaskContextInfo with task_id and session_id, or None if not in a task. + """ + if not is_docket_available(): + return None + + from docket.dependencies import Dependency as DocketDependency + + try: + execution = DocketDependency.execution.get() + # Parse the task key: {session_id}:{task_id}:{task_type}:{component} + from fastmcp.server.tasks.keys import parse_task_key + + key_parts = parse_task_key(execution.key) + return TaskContextInfo( + task_id=key_parts["client_task_id"], + session_id=key_parts["session_id"], + ) + except LookupError: + # Not in worker context + return None + except (ValueError, KeyError): + # Invalid task key format + return None + + +# --- Session registry for background task Context --- + + +_task_sessions: dict[str, weakref.ref[ServerSession]] = {} + + +def register_task_session(session_id: str, session: ServerSession) -> None: + """Register a session for Context access in background tasks. + + Called automatically when a task is submitted to Docket. The session is + stored as a weakref so it doesn't prevent garbage collection when the + client disconnects. + + Args: + session_id: The session identifier + session: The ServerSession instance + """ + _task_sessions[session_id] = weakref.ref(session) + + +def get_task_session(session_id: str) -> ServerSession | None: + """Get a registered session by ID if still alive. + + Args: + session_id: The session identifier + + Returns: + The ServerSession if found and alive, None otherwise + """ + ref = _task_sessions.get(session_id) + if ref is None: + return None + session = ref() + if session is None: + # Session was garbage collected, clean up entry + _task_sessions.pop(session_id, None) + return session + + # --- ContextVars --- _current_server: ContextVar[weakref.ref[FastMCP] | None] = ContextVar( @@ -623,13 +718,52 @@ async def resolve_dependencies( class _CurrentContext(Dependency): # type: ignore[misc] - """Async context manager for Context dependency.""" + """Async context manager for Context dependency. + + In foreground (request) mode: returns the active context from _current_context. + In background (Docket worker) mode: creates a task-aware Context with task_id. + """ + + _context: Context | None = None async def __aenter__(self) -> Context: - return get_context() + from fastmcp.server.context import Context, _current_context + + # Try foreground context first (normal MCP request) + context = _current_context.get() + if context is not None: + return context + + # Check if we're in a Docket worker context + task_info = get_task_context() + if task_info is not None: + # Get session from registry (registered when task was submitted) + session = get_task_session(task_info.session_id) + # Get server from ContextVar + server = get_server() + # Create task-aware Context + self._context = Context( + fastmcp=server, + session=session, + task_id=task_info.task_id, + ) + # Enter the context to set up ContextVars + await self._context.__aenter__() + return self._context + + # Neither foreground nor background context available + raise RuntimeError( + "No active context found. This can happen if:\n" + " - Called outside an MCP request handler\n" + " - Called in a background task before session was registered\n" + "Check `context.request_context` for None before accessing." + ) async def __aexit__(self, *args: object) -> None: - pass + # Clean up if we created a context for background task + if self._context is not None: + await self._context.__aexit__(*args) + self._context = None def CurrentContext() -> Context: diff --git a/src/fastmcp/server/tasks/__init__.py b/src/fastmcp/server/tasks/__init__.py index b3b4a72d4..13ba9e80c 100644 --- a/src/fastmcp/server/tasks/__init__.py +++ b/src/fastmcp/server/tasks/__init__.py @@ -5,6 +5,7 @@ This module implements protocol-level background task execution for MCP servers. from fastmcp.server.tasks.capabilities import get_task_capabilities from fastmcp.server.tasks.config import TaskConfig, TaskMeta, TaskMode +from fastmcp.server.tasks.elicitation import elicit_for_task, handle_task_input from fastmcp.server.tasks.keys import ( build_task_key, get_client_task_id_from_key, @@ -16,7 +17,9 @@ __all__ = [ "TaskMeta", "TaskMode", "build_task_key", + "elicit_for_task", "get_client_task_id_from_key", "get_task_capabilities", + "handle_task_input", "parse_task_key", ] diff --git a/src/fastmcp/server/tasks/config.py b/src/fastmcp/server/tasks/config.py index 1bf0b8ce3..4956a7667 100644 --- a/src/fastmcp/server/tasks/config.py +++ b/src/fastmcp/server/tasks/config.py @@ -7,7 +7,6 @@ handle task-augmented execution as specified in SEP-1686. from __future__ import annotations import inspect -import warnings from collections.abc import Callable from dataclasses import dataclass from datetime import timedelta @@ -136,17 +135,6 @@ class TaskConfig: "Background tasks require async functions." ) - # Warn if function uses Context - it won't be available in workers - from fastmcp.server.context import Context - from fastmcp.utilities.types import find_kwarg_by_type - - context_kwarg = find_kwarg_by_type(fn_to_check, Context) - if context_kwarg: - warnings.warn( - f"'{name}' uses Context but has task execution enabled. " - "Context is not available in background task workers because " - "there is no active MCP session. Consider using Docket dependencies " - "like Progress() instead for worker-compatible functionality.", - UserWarning, - stacklevel=4, - ) + # Note: Context IS now available in background task workers (SEP-1686) + # The wiring in _CurrentContext creates a task-aware Context with task_id + # and session from the registry. No warning needed. diff --git a/src/fastmcp/server/tasks/elicitation.py b/src/fastmcp/server/tasks/elicitation.py new file mode 100644 index 000000000..2fc0bef5f --- /dev/null +++ b/src/fastmcp/server/tasks/elicitation.py @@ -0,0 +1,229 @@ +"""Background task elicitation support (SEP-1686). + +This module provides elicitation capabilities for background tasks running +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 +3. Wait for client to send input via tasks/sendInput +4. Resume task execution with the provided input + +This uses the public MCP SDK APIs where possible, with minimal use of +internal APIs for background task coordination. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import uuid +from typing import TYPE_CHECKING, Any + +import mcp.types +from mcp import ServerSession + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from fastmcp.server.server import FastMCP + + +# Redis key patterns for task elicitation state +ELICIT_REQUEST_KEY = "fastmcp:task:{session_id}:{task_id}:elicit:request" +ELICIT_RESPONSE_KEY = "fastmcp:task:{session_id}:{task_id}:elicit:response" +ELICIT_STATUS_KEY = "fastmcp:task:{session_id}:{task_id}:elicit:status" + +# TTL for elicitation state (1 hour) +ELICIT_TTL_SECONDS = 3600 + + +async def elicit_for_task( + task_id: str, + session: ServerSession, + message: str, + schema: dict[str, Any], + fastmcp: FastMCP, +) -> mcp.types.ElicitResult: + """Send an elicitation request from a background task. + + This function handles the complexity of eliciting user input when running + in a Docket worker context where there's no active MCP request. + + Args: + task_id: The background task ID + session: The MCP ServerSession for this task + message: The message to display to the user + schema: The JSON schema for the expected response + fastmcp: The FastMCP server instance + + Returns: + ElicitResult containing the user's response + + Raises: + RuntimeError: If Docket is not available + McpError: If the elicitation request fails + """ + docket = fastmcp._docket + if docket is None: + raise RuntimeError( + "Background task elicitation requires Docket. " + "Ensure 'fastmcp[tasks]' is installed and the server has task-enabled components." + ) + + # Generate a unique request ID for this elicitation + request_id = str(uuid.uuid4()) + + # Get session ID for Redis key construction + session_id = getattr(session, "_fastmcp_state_prefix", None) + if session_id is None: + # Generate a session ID if not already set + session_id = str(uuid.uuid4()) + session._fastmcp_state_prefix = session_id # type: ignore[attr-defined] + + # Store elicitation request in Redis + request_key = ELICIT_REQUEST_KEY.format(session_id=session_id, task_id=task_id) + response_key = ELICIT_RESPONSE_KEY.format(session_id=session_id, task_id=task_id) + status_key = ELICIT_STATUS_KEY.format(session_id=session_id, task_id=task_id) + + elicit_request = { + "request_id": request_id, + "message": message, + "schema": schema, + } + + async with docket.redis() as redis: + # Store the elicitation request + await redis.set( + docket.key(request_key), + json.dumps(elicit_request), + ex=ELICIT_TTL_SECONDS, + ) + # Set status to "waiting" + await redis.set( + docket.key(status_key), + "waiting", + ex=ELICIT_TTL_SECONDS, + ) + + # Send task status update notification with input_required status + # This follows SEP-1686 for background task status updates + notification = mcp.types.JSONRPCNotification( + jsonrpc="2.0", + method="notifications/tasks/updated", + params={}, + _meta={ # type: ignore[call-arg] + "modelcontextprotocol.io/related-task": { + "taskId": task_id, + "status": "input_required", + "statusMessage": message, + "elicitation": { + "requestId": request_id, + "message": message, + "requestedSchema": schema, + }, + } + }, + ) + + # Send notification (best effort - task status is stored in Redis) + # Log failures for debugging but don't fail the elicitation + try: + await session.send_notification(notification) # type: ignore[arg-type] + except Exception as e: + logger.warning( + "Failed to send input_required notification for task %s: %s", + task_id, + e, + ) + + # Wait for response (poll Redis) + # In a production implementation, this could use Redis pub/sub for lower latency + max_wait_seconds = ELICIT_TTL_SECONDS + poll_interval = 0.5 # seconds + + for _ in range(int(max_wait_seconds / poll_interval)): + async with docket.redis() as redis: + response_data = await redis.get(docket.key(response_key)) + if response_data: + response = json.loads(response_data) + # Clean up Redis keys + await redis.delete( + docket.key(request_key), + docket.key(response_key), + docket.key(status_key), + ) + # Convert to ElicitResult + return mcp.types.ElicitResult( + action=response.get("action", "accept"), + content=response.get("content"), + ) + + await asyncio.sleep(poll_interval) + + # Timeout - treat as cancellation + async with docket.redis() as redis: + await redis.delete( + docket.key(request_key), + docket.key(response_key), + docket.key(status_key), + ) + + return mcp.types.ElicitResult(action="cancel", content=None) + + +async def handle_task_input( + task_id: str, + session_id: str, + action: str, + content: dict[str, Any] | None, + fastmcp: FastMCP, +) -> bool: + """Handle input sent to a background task via tasks/sendInput. + + This is called when a client sends input in response to an elicitation + request from a background task. + + Args: + task_id: The background task ID + session_id: The MCP session ID + action: The elicitation action ("accept", "decline", "cancel") + content: The response content (for "accept" action) + fastmcp: The FastMCP server instance + + Returns: + True if the input was successfully stored, False otherwise + """ + docket = fastmcp._docket + if docket is None: + return False + + response_key = ELICIT_RESPONSE_KEY.format(session_id=session_id, task_id=task_id) + status_key = ELICIT_STATUS_KEY.format(session_id=session_id, task_id=task_id) + + response = { + "action": action, + "content": content, + } + + async with docket.redis() as redis: + # Check if there's a pending elicitation + status = await redis.get(docket.key(status_key)) + if status is None or status.decode("utf-8") != "waiting": + return False + + # Store the response + await redis.set( + docket.key(response_key), + json.dumps(response), + ex=ELICIT_TTL_SECONDS, + ) + # Update status to "responded" + await redis.set( + docket.key(status_key), + "responded", + ex=ELICIT_TTL_SECONDS, + ) + + return True diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index f03dcc1be..02da22148 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/src/fastmcp/server/tasks/handlers.py @@ -101,6 +101,14 @@ async def submit_to_docket( await redis.set(created_at_key, created_at.isoformat(), ex=ttl_seconds) await redis.set(poll_interval_key, str(poll_interval_ms), ex=ttl_seconds) + # Register session for Context access in background workers (SEP-1686) + # This enables elicitation/sampling from background tasks via weakref + # Skip for "internal" sessions (programmatic calls without MCP session) + if session_id != "internal": + from fastmcp.server.dependencies import register_task_session + + 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( diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py new file mode 100644 index 000000000..b778a63b2 --- /dev/null +++ b/tests/server/tasks/test_context_background_task.py @@ -0,0 +1,862 @@ +"""Tests for Context background task support (SEP-1686).""" + +import pytest + +from fastmcp import FastMCP +from fastmcp.server.context import Context +from fastmcp.server.elicitation import AcceptedElicitation +from fastmcp.server.tasks.elicitation import elicit_for_task, handle_task_input + + +class TestContextBackgroundTaskSupport: + """Tests for Context.is_background_task and related functionality.""" + + def test_context_not_background_task_by_default(self): + """Context should not be a background task by default.""" + mcp = FastMCP("test") + ctx = Context(mcp) + assert ctx.is_background_task is False + assert ctx.task_id is None + + def test_context_is_background_task_when_task_id_provided(self): + """Context should be a background task when task_id is provided.""" + mcp = FastMCP("test") + ctx = Context(mcp, task_id="test-task-123") + assert ctx.is_background_task is True + assert ctx.task_id == "test-task-123" + + def test_context_task_id_is_readonly(self): + """task_id should be a read-only property.""" + mcp = FastMCP("test") + ctx = Context(mcp, task_id="test-task-123") + with pytest.raises(AttributeError): + ctx.task_id = "new-id" # type: ignore[misc] + + +class TestContextSessionProperty: + """Tests for Context.session property in different modes.""" + + def test_session_raises_when_no_session_available(self): + """session should raise RuntimeError when no session is available.""" + mcp = FastMCP("test") + ctx = Context(mcp) # No session, not a background task + + with pytest.raises(RuntimeError, match="session is not available"): + _ = ctx.session + + def test_session_uses_stored_session_in_background_task(self): + """session should use _session in background task mode.""" + mcp = FastMCP("test") + + class MockSession: + _fastmcp_state_prefix = "test-session" + + 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): + """session should use _session during on_initialize (no request context).""" + mcp = FastMCP("test") + + class MockSession: + _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" + + ctx._session = MockSession() # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="Docket"): + await ctx.elicit("Need input", str) + + +class TestContextDocumentation: + """Tests to verify Context documentation and API surface.""" + + def test_is_background_task_has_docstring(self): + """is_background_task property should have documentation.""" + assert Context.is_background_task.__doc__ is not None + assert "background task" in Context.is_background_task.__doc__.lower() + + def test_task_id_has_docstring(self): + """task_id property should have documentation.""" + assert Context.task_id.fget.__doc__ is not None + assert "task ID" in Context.task_id.fget.__doc__ + + def test_session_has_docstring(self): + """session property should document background task support.""" + assert Context.session.fget.__doc__ is not None + assert "background task" in Context.session.fget.__doc__.lower() + + +class TestBackgroundTaskElicitationE2E: + """End-to-end tests for background task elicitation (SEP-1686). + + 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. + """ + + 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 + + @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 + 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 + 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 + task_completed = asyncio.Event() + captured_task_id: str | None = None + captured_session_id: str | None = None + captured_is_background: bool | None = None + + @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 + + task_completed.set() + return f"task_id={ctx.task_id}, is_background={ctx.is_background_task}" + + 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) + result = await task.result() + assert "is_background=True" in result.data + + async def test_context_elicit_full_flow_with_mocked_redis(self): + """E2E test with mocked Redis to show complete elicitation flow. + + This test demonstrates what the client sees during background task + elicitation, with a mocked Redis layer to avoid requiring real Redis. + + 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." + """ + 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) + 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 + 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" + + 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 + 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 + success = False + for _ in range(max_attempts): + success = await handle_task_input( + task_id=captured["task_id"], + session_id=captured["session_id"], + action="accept", + content={"value": "Bob"}, + fastmcp=mcp, + ) + if success: + break + await asyncio.sleep(poll_interval_seconds) + + assert success is True, ( + f"handle_task_input should succeed within {max_attempts * poll_interval_seconds}s" + ) + + # 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!"