diff --git a/fastmcp_slim/fastmcp/server/context.py b/fastmcp_slim/fastmcp/server/context.py index 594b59c00..4eeb3bd0a 100644 --- a/fastmcp_slim/fastmcp/server/context.py +++ b/fastmcp_slim/fastmcp/server/context.py @@ -3,7 +3,7 @@ from __future__ import annotations import logging import warnings import weakref -from collections.abc import Callable, Generator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar, Token from dataclasses import dataclass @@ -124,6 +124,32 @@ def _warn_sampling_deprecated() -> None: _current_context: ContextVar[Context | None] = ContextVar("context", default=None) + +#: Hook installed by the tasks extension (``fastmcp-tasks``) so ``ctx.elicit()`` +#: works inside a background-task worker, where there is no live request to +#: carry the elicitation. Core ships no task engine; the extension registers a +#: handler here at construction and ``Context._elicit_for_task`` delegates to it. +#: ``None`` (the default) means no tasks extension is active, so in-task +#: elicitation raises a clear install hint. +_task_elicitation_handler: ( + Callable[[Context, str, dict[str, Any]], Awaitable[mcp_types.ElicitResult]] | None +) = None + + +def set_task_elicitation_handler( + handler: Callable[[Context, str, dict[str, Any]], Awaitable[mcp_types.ElicitResult]] + | None, +) -> None: + """Install (or clear) the in-task elicitation handler. + + Called by the tasks extension so a worker's ``ctx.elicit()`` parks an input + request the client answers via ``tasks/update`` (SEP-2663 poll-based input). + Passing ``None`` restores the default "requires the tasks extension" error. + """ + global _task_elicitation_handler + _task_elicitation_handler = handler + + TransportType = Literal["stdio", "sse", "streamable-http"] _current_transport: ContextVar[TransportType | None] = ContextVar( "transport", default=None @@ -363,6 +389,25 @@ class Context: """ return fastmcp_request_ctx.get() + def client_extension_settings(self, identifier: str) -> dict[str, Any] | None: + """This request's per-request opt-in settings for an MCP extension. + + SEP-2133 extensions negotiate per request: the client repeats its + extension capabilities in each request's ``_meta`` under + ``io.modelcontextprotocol/clientCapabilities`` → ``extensions`` → + ``identifier``. Returns the declared settings dict (possibly empty) when + the extension was opted in for this request, or ``None`` when it was + not (or there is no active request). This bridges an extension's + ``tools/call`` interceptor — which receives a FastMCP ``Context`` — to + the request's declared client capabilities. + """ + rc = self.request_context + if rc is None: + return None + from fastmcp.server.extensions import _extract_client_extension_settings + + return _extract_client_extension_settings(rc.meta, identifier) + def _input_response_params( self, ) -> mcp_types.InputResponseRequestParams | None: @@ -1384,13 +1429,17 @@ class Context: ) # In-task elicitation is provided by the tasks extension (SEP-2663) - # from the `fastmcp-tasks` package. Core no longer ships the SEP-1686 - # push relay this used to call. - raise RuntimeError( - "In-task elicitation requires the tasks extension. Install " - "'fastmcp[tasks]' and register the tasks extension via " - "mcp.add_extension(...)." - ) + # from the `fastmcp-tasks` package, which installs the handler below. + # Core ships no task engine, so without the extension this raises a + # clear install hint rather than reaching a wire the worker lacks. + handler = _task_elicitation_handler + if handler is None: + raise RuntimeError( + "In-task elicitation requires the tasks extension. Install " + "'fastmcp[tasks]' and register the tasks extension via " + "mcp.add_extension(...)." + ) + return await handler(self, message, schema) def _make_state_key(self, key: str) -> str: """Create session-prefixed key for state storage.""" diff --git a/fastmcp_slim/fastmcp/server/dependencies.py b/fastmcp_slim/fastmcp/server/dependencies.py index f0e260cd1..5b34e16ef 100644 --- a/fastmcp_slim/fastmcp/server/dependencies.py +++ b/fastmcp_slim/fastmcp/server/dependencies.py @@ -11,7 +11,7 @@ from __future__ import annotations import importlib.metadata import inspect import weakref -from collections.abc import AsyncGenerator, Callable, Generator, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping from contextlib import AsyncExitStack, asynccontextmanager, contextmanager from contextvars import ContextVar from dataclasses import dataclass @@ -166,6 +166,47 @@ _current_server: ContextVar[weakref.ref[FastMCP] | None] = ContextVar( ) +#: Hook installed by the tasks extension (``fastmcp-tasks``) so a ``ctx: Context`` +#: parameter resolves inside a background-task worker, where there is no +#: foreground request context. Core ships no task engine; the extension +#: registers a factory here that builds and enters a worker ``Context`` (reading +#: the task snapshot restored by the worker). ``_CurrentContext`` falls back to +#: it when no foreground context is active. ``None`` means no tasks extension, +#: so worker context injection is unavailable and the usual "no active context" +#: error applies. +_background_context_factory: Callable[[], Awaitable[Context | None]] | None = None + + +def set_background_context_factory( + factory: Callable[[], Awaitable[Context | None]] | None, +) -> None: + """Install (or clear) the background-task ``Context`` factory. + + The factory returns an already-entered ``Context`` (so ``_current_context`` + is set for cleanup) when called inside a worker, or ``None`` when there is + no task context. Passing ``None`` restores core's no-worker-fallback + behavior. + """ + global _background_context_factory + _background_context_factory = factory + + +#: Hook installed by the tasks extension so ``get_server()`` (and thus +#: ``CurrentFastMCP()``) resolves to the server a mounted task's tool lives on +#: rather than the root that started the worker (#3571). Returns that server +#: inside a worker, or ``None`` outside one. Core has no task engine, so this is +#: ``None`` unless the extension is active. +_worker_server_resolver: Callable[[], FastMCP | None] | None = None + + +def set_worker_server_resolver( + resolver: Callable[[], FastMCP | None] | None, +) -> None: + """Install (or clear) the worker-server resolver used by ``get_server()``.""" + global _worker_server_resolver + _worker_server_resolver = resolver + + # --- Docket availability check --- _DOCKET_AVAILABLE: bool | None = None @@ -360,12 +401,22 @@ def get_context() -> Context: def get_server() -> FastMCP: """Get the current FastMCP server instance directly. + In a background-task worker the tasks extension's resolver is consulted + first, so a mounted-child task resolves to the child server rather than the + root that started the worker (#3571). + Returns: The active FastMCP server Raises: RuntimeError: If no server in context """ + resolver = _worker_server_resolver + if resolver is not None: + worker_server = resolver() + if worker_server is not None: + return worker_server + server_ref = _current_server.get() if server_ref is None: raise RuntimeError("No FastMCP server instance in context") @@ -738,9 +789,20 @@ class _CurrentContext(Dependency["Context"]): if context is not None: return context + # In a background-task worker there is no foreground context; the tasks + # extension installs a factory that builds and enters a worker Context + # from the restored task snapshot. Core has no task engine of its own, + # so this is None unless the extension is active. + factory = _background_context_factory + if factory is not None: + background = await factory() + if background is not None: + return background + raise RuntimeError( "No active context found. This can happen if:\n" " - Called outside an MCP request handler\n" + " - Called in a background task before the context was established\n" "Check `context.request_context` for None before accessing." ) diff --git a/fastmcp_slim/fastmcp/server/extensions.py b/fastmcp_slim/fastmcp/server/extensions.py index e3916f138..91de88d5a 100644 --- a/fastmcp_slim/fastmcp/server/extensions.py +++ b/fastmcp_slim/fastmcp/server/extensions.py @@ -45,8 +45,6 @@ from pydantic import BaseModel from fastmcp.server.dependencies import _lift_meta, bind_request_context if TYPE_CHECKING: - import mcp_types - from fastmcp.server.context import Context from fastmcp.server.server import FastMCP from fastmcp.tools.base import ToolResult @@ -58,8 +56,10 @@ __all__ = [ ] # What an extension's tools/call interceptor observes and may produce: the tool -# result, or the claimed CreateTaskResult shape when the call is run as a task. -ToolCallOutcome: TypeAlias = "ToolResult | mcp_types.CreateTaskResult" +# result, or an extension-defined wire result model (a `BaseModel` the runner +# serializes) when the call is short-circuited — e.g. the tasks extension's +# CreateTaskResult. Core does not interpret the extension's result shape. +ToolCallOutcome: TypeAlias = "ToolResult | BaseModel" # A method handler receives the SDK request context plus validated params and # returns a bare result model (the runner serializes it). diff --git a/fastmcp_slim/fastmcp/server/mixins/lifespan.py b/fastmcp_slim/fastmcp/server/mixins/lifespan.py index f460e1872..79699da00 100644 --- a/fastmcp_slim/fastmcp/server/mixins/lifespan.py +++ b/fastmcp_slim/fastmcp/server/mixins/lifespan.py @@ -118,6 +118,13 @@ class LifespanMixin: """ from fastmcp.utilities.tasks import TASKS_EXTENSION_ID + # A mounted child defers to the root, which owns the extension and whose + # aggregated get_tasks() already covers this child's task tools — the + # same root-deferral the extension lifespan uses. Validating here would + # fail a child that legitimately relies on the root's registration. + if _lifespan_root_active.get(): + return + if TASKS_EXTENSION_ID in self._extensions: return diff --git a/fastmcp_slim/fastmcp/server/mixins/mcp_operations.py b/fastmcp_slim/fastmcp/server/mixins/mcp_operations.py index 3bc8e96ba..af85dd08c 100644 --- a/fastmcp_slim/fastmcp/server/mixins/mcp_operations.py +++ b/fastmcp_slim/fastmcp/server/mixins/mcp_operations.py @@ -20,6 +20,7 @@ from mcp_types import ( SetLevelRequestParams, ) from mcp_types.version import MODERN_PROTOCOL_VERSIONS +from pydantic import BaseModel from fastmcp.exceptions import ( DisabledError, @@ -29,7 +30,7 @@ from fastmcp.exceptions import ( ) from fastmcp.server.completions import CompletionValues, normalize_completion from fastmcp.server.dependencies import bind_request_context, extract_version_spec -from fastmcp.tools.base import InputRequiredToolResult +from fastmcp.tools.base import InputRequiredToolResult, ToolResult from fastmcp.utilities.async_utils import ( call_sync_fn_in_threadpool, is_coroutine_function, @@ -216,7 +217,7 @@ class MCPOperationsMixin: self: FastMCP, ctx: ServerRequestContext, params: CallToolRequestParams, - ) -> mcp_types.CallToolResult | mcp_types.InputRequiredResult: + ) -> mcp_types.CallToolResult | mcp_types.InputRequiredResult | BaseModel: """Handle MCP 'tools/call' requests. A guard tool (SEP-2322 multi-round-trip) requests client input by @@ -263,6 +264,14 @@ class MCPOperationsMixin: is_error=True, ) + if not isinstance(result, ToolResult): + # An extension's tools/call interceptor produced a non-ToolResult + # wire result — the tasks extension's CreateTaskResult when it ran + # the call as a task. Core does not interpret extension result + # shapes; hand it straight to the runner, which serializes it for + # the negotiated protocol version. + return result + if isinstance(result, InputRequiredToolResult): # A guard tool requested client input (SEP-2322). The # multi-round-trip result type only exists at 2026-07-28; on an diff --git a/fastmcp_slim/fastmcp/server/server.py b/fastmcp_slim/fastmcp/server/server.py index 1308013ad..15e03c5f2 100644 --- a/fastmcp_slim/fastmcp/server/server.py +++ b/fastmcp_slim/fastmcp/server/server.py @@ -654,7 +654,15 @@ class FastMCP( can reach it), its method bindings are wired onto the low-level server, and it is recorded for capability advertisement, interception, and lifespan entry. Registering two extensions with the same identifier is - an error. + an error, as is registering after the server's lifespan has started — + the extension's lifespan could no longer run, leaving it silently + half-active. + + Extensions are served by the server they are registered on. A mounted + child's extensions do not propagate to the root: the root serves the + wire, so only root-registered extensions advertise capabilities and + answer methods (matching the lifespan, which also defers to the root). + Register extensions on the server you run. """ from fastmcp.server.extensions import ( build_method_handler, @@ -669,6 +677,12 @@ class FastMCP( f"An extension with identifier {extension.identifier!r} is " "already registered." ) + if self._lifespan_result_set: + raise RuntimeError( + f"Cannot register extension {extension.identifier!r}: the " + "server's lifespan has already started, so the extension's " + "lifespan would never run. Register extensions before serving." + ) extension._bind(self) for binding in extension.methods(): diff --git a/fastmcp_tasks/fastmcp_tasks/__init__.py b/fastmcp_tasks/fastmcp_tasks/__init__.py index 9880ae3b6..32c500f0e 100644 --- a/fastmcp_tasks/fastmcp_tasks/__init__.py +++ b/fastmcp_tasks/fastmcp_tasks/__init__.py @@ -2,9 +2,11 @@ from importlib.metadata import PackageNotFoundError, version +from fastmcp_tasks.extension import TasksExtension + try: __version__ = version("fastmcp-tasks") except PackageNotFoundError: __version__ = "0.0.0" -__all__ = ["__version__"] +__all__ = ["TasksExtension", "__version__"] diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/__init__.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/__init__.py deleted file mode 100644 index 57bbff5ee..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""SEP-1686 wire layer, moved intact and awaiting Phase 3 adaptation. - -Every module in this subpackage is the original SEP-1686-shaped wire code: -the four CRUD request handlers (`requests.py`), the task-submission handler -(`handlers.py`), the Docket-subscription status relay (`subscriptions.py`), the -Redis push relay for elicitation (`elicitation.py`, `notifications.py`), the -capability declaration (`capabilities.py`), and the mode-routing dispatcher -(`routing.py`). - -It is disconnected from core — nothing wires these handlers onto a server after -Phase 2. Phase 3 adapts this code in place to the SEP-2663 `tasks/get|update|cancel` -shape under its ported tests. Do not "improve" it here; the point of keeping it is -that it embodies operational lessons the rewrite must preserve. -""" diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/capabilities.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/capabilities.py deleted file mode 100644 index 42367668c..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/capabilities.py +++ /dev/null @@ -1,45 +0,0 @@ -"""SEP-1686 task capabilities declaration.""" - -from mcp_types import ( - ServerTasksCapability, - ServerTasksRequestsCapability, - TasksCallCapability, - TasksCancelCapability, - TasksListCapability, - TasksToolsCapability, -) - - -def get_task_capabilities() -> ServerTasksCapability | None: - """Return the SEP-1686 task capabilities. - - Returns task capabilities as a first-class ServerCapabilities field, - declaring support for list, cancel, and request operations per SEP-1686. - - Returns None if a compatible pydocket is not installed (no task support). - Uses the canonical ``is_docket_available()`` check so that capability - advertisement and handler registration stay in sync — otherwise a server - with an old transitive pydocket would advertise task support and then - return "method not found" when clients invoked it. - - Only tools are advertised as task-capable. In the SDK v2 b1 wire types, - ``ReadResourceRequestParams`` / ``GetPromptRequestParams`` carry no ``task`` - field (sdk-feedback #3), so resource/prompt task submissions are not - wire-expressible and always graceful-degrade to synchronous execution. - Advertising ``prompts``/``resources`` task support would mislead - capability-discovering clients into sending task-augmented reads/gets that - silently run synchronously. Restore them here once the SDK adds task - metadata to those request params. - """ - from fastmcp_tasks.dependencies import is_docket_available - - if not is_docket_available(): - return None - - return ServerTasksCapability( - list=TasksListCapability(), - cancel=TasksCancelCapability(), - requests=ServerTasksRequestsCapability( - tools=TasksToolsCapability(call=TasksCallCapability()), - ), - ) diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/elicitation.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/elicitation.py deleted file mode 100644 index 95798e1f1..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/elicitation.py +++ /dev/null @@ -1,347 +0,0 @@ -"""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/status 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 json -import logging -import uuid -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any - -import mcp_types -from mcp import ServerSession - -from fastmcp_tasks._legacy_wire.notifications import push_notification -from fastmcp_tasks.context import get_task_context, get_task_session_id -from fastmcp_tasks.keys import task_redis_prefix - -logger = logging.getLogger(__name__) - -if TYPE_CHECKING: - from fastmcp.server.server import FastMCP - - -# TTL for elicitation state (1 hour) -ELICIT_TTL_SECONDS = 3600 - - -def _elicit_keys(task_scope: str | None, task_id: str) -> tuple[str, str, str]: - """Build (request, response, status) Redis keys for a task's elicitation.""" - prefix = f"{task_redis_prefix(task_scope)}:{task_id}:elicit" - return f"{prefix}:request", f"{prefix}:response", f"{prefix}:status" - - -async def elicit_for_task( - task_id: str, - session: ServerSession | None, - 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()) - - task_context = get_task_context() - if task_context is not None: - task_scope = task_context.task_scope - # Prefer the live session's cached ID (always available in-process), - # fall back to the snapshot for distributed workers. - session_id = ( - getattr(session, "_fastmcp_state_prefix", None) or get_task_session_id() - ) - else: - raise RuntimeError( - "Cannot determine task scope for elicitation. " - "This typically means elicit_for_task() was called outside a Docket worker context." - ) - - # Store elicitation request in Redis - request_key, response_key, status_key = _elicit_keys(task_scope, 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. - # 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/status", - "params": { - "taskId": task_id, - "status": "input_required", - "statusMessage": message, - "createdAt": timestamp, - "lastUpdatedAt": timestamp, - "ttl": ELICIT_TTL_SECONDS * 1000, - }, - "_meta": { - "io.modelcontextprotocol/related-task": { - "taskId": task_id, - "status": "input_required", - "statusMessage": message, - "task_scope": task_scope, - "elicitation": { - "requestId": request_id, - "message": message, - "requestedSchema": schema, - }, - } - }, - } - - if session_id is None: - logger.warning( - "No session_id available for task %s, cannot deliver elicitation notification", - task_id, - ) - return mcp_types.ElicitResult(action="cancel", content=None) - - try: - await push_notification(session_id, notification_dict, docket) - except Exception as e: - # Fail fast: if notification can't be queued, client won't know to respond - # Return cancel immediately rather than waiting for 1-hour timeout - logger.warning( - "Failed to queue input_required notification for task %s, cancelling elicitation: %s", - task_id, - e, - ) - # Best-effort cleanup - try: - async with docket.redis() as redis: - await redis.delete( - docket.key(request_key), - docket.key(status_key), - ) - except Exception: - pass # Keys will expire via TTL - return mcp_types.ElicitResult(action="cancel", content=None) - - # Wait for response using BLPOP (blocking pop) - # This is much more efficient than polling - single Redis round-trip - # that blocks until a response is pushed, vs 7,200 round-trips/hour with polling - max_wait_seconds = ELICIT_TTL_SECONDS - - try: - 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, - ) - - if result: - # result is (key, value) tuple - _key, response_data = result - response = json.loads(response_data) - - # Clean up Redis keys - await redis.delete( - docket.key(request_key), - docket.key(status_key), - ) - - # Convert to ElicitResult - return mcp_types.ElicitResult( - action=response.get("action", "accept"), - content=response.get("content"), - ) - except Exception as e: - logger.warning( - "BLPOP failed for task %s elicitation, falling back to cancel: %s", - task_id, - e, - ) - - # Timeout or error - treat as cancellation - # Best-effort cleanup - if Redis is unavailable, keys will expire via TTL - try: - async with docket.redis() as redis: - await redis.delete( - docket.key(request_key), - docket.key(response_key), - docket.key(status_key), - ) - except Exception as cleanup_error: - logger.debug( - "Failed to clean up elicitation keys for task %s (will expire via TTL): %s", - task_id, - cleanup_error, - ) - - return mcp_types.ElicitResult(action="cancel", content=None) - - -async def relay_elicitation( - session: ServerSession, - task_scope: str | None, - task_id: str, - elicitation: dict[str, Any], - fastmcp: FastMCP, -) -> None: - """Relay elicitation from a background task worker to the client. - - Called by the notification subscriber when it detects an input_required - notification with elicitation metadata. Sends a standard elicitation/create - request to the client session, then uses handle_task_input() to push the - response to Redis so the blocked worker can resume. - - Args: - session: MCP ServerSession - task_scope: Authorization scope for Redis key construction - task_id: Background task ID - elicitation: Elicitation metadata (message, requestedSchema) - fastmcp: FastMCP server instance - """ - try: - result = await session.elicit( - message=elicitation["message"], - requested_schema=elicitation["requestedSchema"], - ) - await handle_task_input( - task_id=task_id, - task_scope=task_scope, - action=result.action, - content=result.content, - fastmcp=fastmcp, - ) - logger.debug( - "Relayed elicitation response for task %s (action=%s)", - task_id, - result.action, - ) - except Exception as e: - logger.warning("Failed to relay elicitation for task %s: %s", task_id, e) - # Push a cancel response so the worker's BLPOP doesn't block forever - success = await handle_task_input( - task_id=task_id, - task_scope=task_scope, - action="cancel", - content=None, - fastmcp=fastmcp, - ) - if not success: - logger.warning( - "Failed to push cancel response for task %s " - "(worker may block until TTL)", - task_id, - ) - - -async def handle_task_input( - task_id: str, - task_scope: str | None, - 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 - task_scope: Authorization scope for Redis key construction - 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, status_key = _elicit_keys(task_scope, 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 - - # 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( - docket.key(response_key), - json.dumps(response), - ) - # Set TTL on the response list (in case BLPOP doesn't consume it) - await redis.expire(docket.key(response_key), ELICIT_TTL_SECONDS) - - # Update status to "responded" - await redis.set( - docket.key(status_key), - "responded", - ex=ELICIT_TTL_SECONDS, - ) - - return True diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/handlers.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/handlers.py deleted file mode 100644 index 2ff531c66..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/handlers.py +++ /dev/null @@ -1,266 +0,0 @@ -"""SEP-1686 task execution handlers. - -Handles queuing tool/prompt/resource executions to Docket as background tasks. -""" - -from __future__ import annotations - -import asyncio -import uuid -from contextlib import suppress -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Literal - -import mcp_types -from mcp.shared.exceptions import MCPError -from mcp_types import INTERNAL_ERROR - -from fastmcp.server.dependencies import get_context -from fastmcp.tools.function_tool import _strict_input_validation -from fastmcp.utilities.logging import get_logger -from fastmcp.utilities.tasks import TaskMeta -from fastmcp_tasks.components import add_component_to_docket, coerce_task_arguments -from fastmcp_tasks.context import ( - TaskContextSnapshot, - get_task_scope, - register_task_server, - register_task_session, -) -from fastmcp_tasks.dependencies import _current_docket -from fastmcp_tasks.keys import build_task_key, task_redis_prefix - -if TYPE_CHECKING: - from fastmcp.prompts.base import Prompt - from fastmcp.resources.base import Resource - from fastmcp.resources.template import ResourceTemplate - from fastmcp.tools.base import Tool - -logger = get_logger(__name__) - -# Redis mapping TTL buffer: Add 15 minutes to Docket's execution_ttl -TASK_MAPPING_TTL_BUFFER_SECONDS = 15 * 60 - - -async def submit_to_docket( - task_type: Literal["tool", "resource", "template", "prompt"], - key: str, - component: Tool | Resource | ResourceTemplate | Prompt, - arguments: dict[str, Any] | None = None, - task_meta: TaskMeta | None = None, -) -> mcp_types.CreateTaskResult: - """Submit any component to Docket for background execution (SEP-1686). - - Unified handler for all component types. Called by component's internal - methods (_run, _read, _render) when task metadata is present and mode allows. - - Queues the component's method to Docket, stores raw return values, - and converts to MCP types on retrieval. - - Args: - task_type: Component type for task key construction - key: The component key as seen by MCP layer (with namespace prefix) - component: The component instance (Tool, Resource, ResourceTemplate, Prompt) - arguments: Arguments/params (None for Resource which has no args) - task_meta: Task execution metadata. If task_meta.ttl is provided, it - overrides the server default (docket.execution_ttl). - - Returns: - CreateTaskResult: Task stub with proper Task object - """ - # Validate and coerce arguments before creating any task state. A failure - # here must surface before the Redis metadata and initial "working" - # notification below are written, otherwise an invalid input would orphan a - # task the client has already observed (#4349). - # - # Honor the server's strict_input_validation setting so a strict tool - # rejects lax coercions (e.g. {"n": "1"} for n: int) at submission just as - # it does on the synchronous call path — otherwise task=True would bypass - # strict validation entirely. - if arguments is not None: - arguments = coerce_task_arguments( - component, arguments, strict=_strict_input_validation() - ) - - # Generate server-side task ID per SEP-1686 final spec (line 375-377) - # Server MUST generate task IDs, clients no longer provide them - server_task_id = str(uuid.uuid4()) - - # Record creation timestamp per SEP-1686 final spec (line 430). SDK v2 - # types `Task.created_at` / `TaskStatusNotificationParams.created_at` as ISO - # strings, so carry a serialized copy for wire-crossing models. - created_at = datetime.now(timezone.utc) - created_at_iso = created_at.isoformat() - - ctx = get_context() - - # Authorization scope for task isolation (auth identity, or None for anonymous) - task_scope = get_task_scope() - - # Transport session ID for notification delivery - try: - session_id = ctx.session_id - except RuntimeError: - session_id = None - - # Try the server's own Docket first; fall back to the ContextVar for - # mounted children (whose parent server owns the Docket instance). - docket = ctx.fastmcp._docket or _current_docket.get() - if docket is None: - raise MCPError( - code=INTERNAL_ERROR, - message="Background tasks require a running FastMCP server context", - ) - - # Register the current server so background workers resolve - # CurrentFastMCP() / ctx.fastmcp to the correct (child) server - # for mounted tasks. At this point ctx.fastmcp is the child because - # we're inside the child's call_tool dispatch. - register_task_server(server_task_id, ctx.fastmcp) - - # Build full task key with embedded metadata - task_key = build_task_key(task_scope, server_task_id, task_type, key) - - # Determine TTL: use task_meta.ttl if provided, else docket default - if task_meta is not None and task_meta.ttl is not None: - ttl_ms = task_meta.ttl - else: - ttl_ms = int(docket.execution_ttl.total_seconds() * 1000) - ttl_seconds = int(ttl_ms / 1000) + TASK_MAPPING_TTL_BUFFER_SECONDS - - # Store task metadata in Redis for protocol handlers - prefix = task_redis_prefix(task_scope) - task_meta_key = docket.key(f"{prefix}:{server_task_id}") - created_at_key = docket.key(f"{prefix}:{server_task_id}:created_at") - poll_interval_key = docket.key(f"{prefix}:{server_task_id}:poll_interval") - poll_interval_ms = int(component.task_config.poll_interval.total_seconds() * 1000) - - # Snapshot all context (access token, headers, origin request ID, - # and session_id for notification delivery in background workers) - snapshot = TaskContextSnapshot.capture() - - async with docket.redis() as redis: - await redis.set(task_meta_key, task_key, ex=ttl_seconds) - 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) - - await snapshot.save(docket, task_scope, server_task_id, ttl_seconds) - - # Register session for Context access in background workers (SEP-1686) - # This enables elicitation/sampling from background tasks via weakref - # Skip when there is no session (programmatic calls without MCP session) - if session_id is not None: - register_task_session(session_id, ctx.session) - - # 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_iso, - "lastUpdatedAt": created_at_iso, - "ttl": ttl_ms, - "pollInterval": poll_interval_ms, - }, - "_meta": { - "io.modelcontextprotocol/related-task": { - "taskId": server_task_id, - } - }, - } - ) - # SDK v2: `ServerNotification` is a union type, not a wrapper class; - # `send_notification` takes the bare notification model directly. - with suppress(Exception): - # Don't let notification failures break task creation - await ctx.session.send_notification(notification) # type: ignore[arg-type] # ty:ignore[invalid-argument-type] - - # Queue function to Docket by key (result storage via execution_ttl) - # Use component.add_to_docket() which handles calling conventions - # `fn_key` is the function lookup key (e.g., "child_multiply") - # `task_key` is the task result key (e.g., "fastmcp:task:{task_scope}:{task_id}:tool:child_multiply") - # Resources don't take arguments; tools/prompts/templates always pass arguments (even if None/empty) - if task_type == "resource": - await add_component_to_docket( - component, docket, None, fn_key=key, task_key=task_key - ) - else: - await add_component_to_docket( - component, docket, arguments, fn_key=key, task_key=task_key - ) - - # Spawn subscription task to send status notifications (SEP-1686 optional feature). - # SDK v2 constructs a ServerSession per request and exposes no per-connection - # task group, so the subscription runs as a standalone asyncio task that - # outlives the submitting request; it is cancelled when the connection closes. - # Deferred: subscriptions and notifications depend on docket at import time - from fastmcp_tasks._legacy_wire.subscriptions import subscribe_to_task_updates - - subscription_task = asyncio.create_task( - subscribe_to_task_updates( - server_task_id, - task_key, - ctx.session, - docket, - poll_interval_ms, - ), - name=f"task-subscription-{server_task_id[:8]}", - ) - connection = getattr(ctx.session, "_connection", None) - if connection is not None: - - async def _cancel_subscription() -> None: - if not subscription_task.done(): - subscription_task.cancel() - with suppress(asyncio.CancelledError): - await subscription_task - - connection.exit_stack.push_async_callback(_cancel_subscription) - - # Deferred: notifications depends on docket at import time - from fastmcp_tasks._legacy_wire.notifications import ( - ensure_subscriber_running, - stop_subscriber, - ) - - if session_id is not None: - try: - await ensure_subscriber_running( - session_id, ctx.session, docket, ctx.fastmcp - ) - - # Register cleanup callback on connection exit (once per session). - # SDK v2 constructs ServerSession per request, so the stable - # per-connection lifecycle hook lives on the underlying Connection - # (`connection.exit_stack`), not the session. The registration flag - # is likewise stashed on the connection's `state` so it survives - # across requests. - connection = getattr(ctx.session, "_connection", None) - if connection is not None and not connection.state.get( - "_notification_cleanup_registered" - ): - - async def _cleanup_subscriber() -> None: - await stop_subscriber(session_id) # type: ignore[arg-type] - - connection.exit_stack.push_async_callback(_cleanup_subscriber) - connection.state["_notification_cleanup_registered"] = True - except Exception as e: - # Non-fatal: elicitation will still work via polling fallback - logger.debug("Failed to start notification subscriber: %s", e) - - # Return CreateTaskResult with proper Task object - # Tasks MUST begin in "working" status per SEP-1686 final spec (line 381) - return mcp_types.CreateTaskResult( - task=mcp_types.Task( - task_id=server_task_id, - status="working", - created_at=created_at_iso, - last_updated_at=created_at_iso, - ttl=ttl_ms, - poll_interval=poll_interval_ms, - ) - ) diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/notifications.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/notifications.py deleted file mode 100644 index 1affd89af..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/notifications.py +++ /dev/null @@ -1,312 +0,0 @@ -"""Distributed notification queue for background task events (SEP-1686). - -Enables distributed Docket workers to send MCP notifications to clients -without holding session references. Workers push to a Redis queue, -the MCP server process subscribes and forwards to the client's session. - -Pattern: Fire-and-forward with retry -- One queue per session_id -- LPUSH/BRPOP for reliable ordered delivery -- Retry up to 3 times on delivery failure, then discard -- TTL-based expiration for stale messages - -Note: Docket's execution.subscribe() handles task state/progress events via -Redis Pub/Sub. This module handles elicitation-specific notifications that -require reliable delivery (input_required prompts, cancel signals). -""" - -from __future__ import annotations - -import asyncio -import json -import logging -import weakref -from contextlib import suppress -from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any - -import mcp_types - -if TYPE_CHECKING: - from docket import Docket - from mcp.server.session import ServerSession - - from fastmcp.server.server import FastMCP - -logger = logging.getLogger(__name__) - -# Redis key patterns -NOTIFICATION_QUEUE_KEY = "fastmcp:notifications:{session_id}" -NOTIFICATION_ACTIVE_KEY = "fastmcp:notifications:{session_id}:active" - -# Configuration -NOTIFICATION_TTL_SECONDS = 300 # 5 minute message TTL (elicitation response window) -MAX_DELIVERY_ATTEMPTS = 3 # Retry failed deliveries before discarding -SUBSCRIBER_TIMEOUT_SECONDS = 30 # BRPOP timeout (also heartbeat interval) - - -async def push_notification( - session_id: str, - notification: dict[str, Any], - docket: Docket, -) -> None: - """Push notification to session's queue (called from Docket worker). - - Used for elicitation-specific notifications (input_required, cancel) - that need reliable delivery across distributed processes. - - Args: - session_id: Target session's identifier - notification: MCP notification dict (method, params, _meta) - docket: Docket instance for Redis access - """ - key = docket.key(NOTIFICATION_QUEUE_KEY.format(session_id=session_id)) - message = json.dumps( - { - "notification": notification, - "attempt": 0, - "enqueued_at": datetime.now(timezone.utc).isoformat(), - } - ) - async with docket.redis() as redis: - await redis.lpush(key, message) - await redis.expire(key, NOTIFICATION_TTL_SECONDS) - - -async def notification_subscriber_loop( - session_id: str, - session: ServerSession, - docket: Docket, - fastmcp: FastMCP, -) -> None: - """Subscribe to notification queue and forward to session. - - Runs in the MCP server process. Bridges distributed workers to clients. - - This loop: - 1. Maintains a heartbeat (active subscriber marker for debugging) - 2. Blocks on BRPOP waiting for notifications - 3. Forwards notifications to the client's session - 4. Retries failed deliveries, then discards (no dead-letter queue) - - Args: - session_id: Session identifier to subscribe to - session: MCP ServerSession for sending notifications - docket: Docket instance for Redis access - fastmcp: FastMCP server instance (for elicitation relay) - """ - queue_key = docket.key(NOTIFICATION_QUEUE_KEY.format(session_id=session_id)) - active_key = docket.key(NOTIFICATION_ACTIVE_KEY.format(session_id=session_id)) - - logger.debug("Starting notification subscriber for session %s", session_id) - - while True: - try: - async with docket.redis() as redis: - # Heartbeat: mark subscriber as active (for distributed debugging) - await redis.set(active_key, "1", ex=SUBSCRIBER_TIMEOUT_SECONDS * 2) - - # 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 - ) - if not result: - continue # Timeout - refresh heartbeat and retry - - _, message_bytes = result - message = json.loads(message_bytes) - notification_dict = message["notification"] - attempt = message.get("attempt", 0) - - try: - # Reconstruct and send MCP notification - await _send_mcp_notification( - session, notification_dict, session_id, docket, fastmcp - ) - logger.debug( - "Delivered notification to session %s (attempt %d)", - session_id, - attempt + 1, - ) - except Exception as send_error: - # Delivery failed - retry or discard - if attempt < MAX_DELIVERY_ATTEMPTS - 1: - # 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)) - logger.debug( - "Requeued notification for session %s (attempt %d): %s", - session_id, - attempt + 2, - send_error, - ) - else: - # Discard after max attempts (session likely disconnected) - logger.warning( - "Discarding notification for session %s after %d attempts: %s", - session_id, - MAX_DELIVERY_ATTEMPTS, - send_error, - ) - - except asyncio.CancelledError: - # Graceful shutdown - leave pending messages in queue for reconnect - logger.debug("Notification subscriber cancelled for session %s", session_id) - break - except Exception as e: - logger.debug( - "Notification subscriber error for session %s: %s", session_id, e - ) - await asyncio.sleep(1) # Backoff on error - - -async def _send_mcp_notification( - session: ServerSession, - notification_dict: dict[str, Any], - session_id: str, - docket: Docket, - fastmcp: FastMCP, -) -> None: - """Reconstruct MCP notification from dict and send to session. - - For input_required notifications with elicitation metadata, also sends - a standard elicitation/create request to the client and relays the - response back to the worker via Redis. - - Args: - session: MCP ServerSession - notification_dict: Notification as dict (method, params, _meta) - session_id: Session identifier (for elicitation relay) - docket: Docket instance (for notification delivery) - fastmcp: FastMCP server instance (for elicitation relay) - """ - method = notification_dict.get("method", "notifications/tasks/status") - if method != "notifications/tasks/status": - raise ValueError(f"Unsupported notification method for subscriber: {method}") - - # SDK v2: a notification's `_meta` lives on its params (`params._meta`), not - # at the notification envelope level, so nest it under params before parsing. - params_dict = dict(notification_dict.get("params", {})) - meta_dict = notification_dict.get("_meta") - if meta_dict is not None: - params_dict["_meta"] = meta_dict - notification = mcp_types.TaskStatusNotification.model_validate( - { - "method": "notifications/tasks/status", - "params": params_dict, - } - ) - # SDK v2: `ServerNotification` is a union type; send the bare model. - await session.send_notification(notification) # type: ignore[arg-type] # ty:ignore[invalid-argument-type] - - # If this is an input_required notification with elicitation metadata, - # relay the elicitation to the client via standard elicitation/create - params = notification_dict.get("params", {}) - if params.get("status") == "input_required": - meta = notification_dict.get("_meta", {}) - related_task = meta.get("io.modelcontextprotocol/related-task", {}) - elicitation = related_task.get("elicitation") - if elicitation: - task_id = params.get("taskId") - if not task_id: - logger.warning( - "input_required notification missing taskId, skipping relay" - ) - return - if "task_scope" not in related_task: - logger.warning( - "input_required notification for task %s missing task_scope " - "metadata, skipping elicitation relay", - task_id, - ) - return - task_scope = related_task["task_scope"] - from fastmcp_tasks._legacy_wire.elicitation import relay_elicitation - - task = asyncio.create_task( - relay_elicitation(session, task_scope, task_id, elicitation, fastmcp), - name=f"elicitation-relay-{task_id[:8]}", - ) - _background_tasks.add(task) - task.add_done_callback(_background_tasks.discard) - - -# ============================================================================= -# Subscriber Management -# ============================================================================= - -# Strong references to fire-and-forget relay tasks (prevent GC mid-flight) -_background_tasks: set[asyncio.Task[None]] = set() - -# Registry of active subscribers per session (prevents duplicates) -# Uses weakref to session to detect disconnects -_active_subscribers: dict[ - str, tuple[asyncio.Task[None], weakref.ref[ServerSession]] -] = {} - - -async def ensure_subscriber_running( - session_id: str, - session: ServerSession, - docket: Docket, - fastmcp: FastMCP, -) -> None: - """Start notification subscriber if not already running (idempotent). - - Subscriber is created on first task submission and cleaned up on disconnect. - Safe to call multiple times for the same session. - - Args: - session_id: Session identifier - session: MCP ServerSession - docket: Docket instance - fastmcp: FastMCP server instance (for elicitation relay) - """ - # Check if subscriber already running for this session - if session_id in _active_subscribers: - task, session_ref = _active_subscribers[session_id] - # Check if task is still running AND session is still alive - if not task.done() and session_ref() is not None: - return # Already running - - # Task finished or session dead - clean up - if not task.done(): - task.cancel() - with suppress(asyncio.CancelledError): - await task - del _active_subscribers[session_id] - - # Start new subscriber task - task = asyncio.create_task( - notification_subscriber_loop(session_id, session, docket, fastmcp), - name=f"notification-subscriber-{session_id[:8]}", - ) - _active_subscribers[session_id] = (task, weakref.ref(session)) - logger.debug("Started notification subscriber for session %s", session_id) - - -async def stop_subscriber(session_id: str) -> None: - """Stop notification subscriber for a session. - - Called when session disconnects. Pending messages remain in queue - for delivery if client reconnects (with TTL expiration). - - Args: - session_id: Session identifier - """ - if session_id not in _active_subscribers: - return - - task, _ = _active_subscribers.pop(session_id) - if not task.done(): - task.cancel() - with suppress(asyncio.CancelledError): - await task - logger.debug("Stopped notification subscriber for session %s", session_id) - - -def get_subscriber_count() -> int: - """Get number of active subscribers (for monitoring).""" - return len(_active_subscribers) diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/requests.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/requests.py deleted file mode 100644 index 5f1f72a0f..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/requests.py +++ /dev/null @@ -1,469 +0,0 @@ -"""SEP-1686 task request handlers. - -Handles MCP task protocol requests: tasks/get, tasks/result, tasks/list, tasks/cancel. -These handlers query and manage existing tasks (contrast with handlers.py which creates tasks). - -This module requires fastmcp[tasks] (pydocket). It is only imported when docket is available. -""" - -from __future__ import annotations - -from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Literal - -import mcp_types -from docket.execution import ExecutionState -from mcp.shared.exceptions import MCPError -from mcp_types import ( - INTERNAL_ERROR, - INVALID_PARAMS, - CancelTaskResult, - GetTaskResult, - ListTasksResult, -) - -import fastmcp.server.context -from fastmcp.exceptions import NotFoundError -from fastmcp.prompts.base import Prompt -from fastmcp.resources.base import Resource -from fastmcp.resources.template import ResourceTemplate -from fastmcp.tools.base import InputRequiredToolResult, Tool -from fastmcp.utilities.tasks import DEFAULT_POLL_INTERVAL_MS, DEFAULT_TTL_MS -from fastmcp.utilities.versions import VersionSpec -from fastmcp_tasks.context import get_task_scope -from fastmcp_tasks.keys import parse_task_key, task_redis_prefix - -if TYPE_CHECKING: - from fastmcp.server.server import FastMCP - - -# Map Docket execution states to MCP task status strings -# Per SEP-1686 final spec (line 381): tasks MUST begin in "working" status -DOCKET_TO_MCP_STATE: dict[ExecutionState, str] = { - ExecutionState.SCHEDULED: "working", # Initial state per spec - ExecutionState.QUEUED: "working", # Initial state per spec - ExecutionState.RUNNING: "working", - ExecutionState.COMPLETED: "completed", - ExecutionState.FAILED: "failed", - ExecutionState.CANCELLED: "cancelled", -} - - -def _normalize_iso_timestamp(stored: str | None) -> str: - """Return an ISO 8601 timestamp string for a Task's createdAt/lastUpdatedAt. - - The v2 Task model types these fields as ISO 8601 strings. `stored` is the - value read from Redis (already an ISO string) or None; either way this - returns a valid ISO string, falling back to the current UTC time. - """ - if stored: - try: - return datetime.fromisoformat(stored.replace("Z", "+00:00")).isoformat() - except (ValueError, AttributeError): - pass - return datetime.now(timezone.utc).isoformat() - - -def _parse_key_version(key_suffix: str) -> tuple[str, str | None]: - """Parse a key suffix into (name_or_uri, version). - - Keys always contain @ as a version delimiter (sentinel pattern): - - "add@1.0" → ("add", "1.0") # versioned - - "add@" → ("add", None) # unversioned - - "user@example.com@1.0" → ("user@example.com", "1.0") # @ in URI - - Uses rsplit to split on the LAST @ which is always the version delimiter. - Falls back to treating the whole string as the name if @ is not present - (for backwards compatibility with legacy task keys). - """ - if "@" not in key_suffix: - # Legacy key without version sentinel - treat as unversioned - return key_suffix, None - name_or_uri, version = key_suffix.rsplit("@", 1) - return name_or_uri, version if version else None - - -async def _lookup_task_execution( - docket: Any, - task_scope: str | None, - client_task_id: str, -) -> tuple[Any, str | None, int]: - """Look up task execution and metadata from Redis. - - Consolidates the common pattern of fetching task metadata from Redis, - validating it exists, and retrieving the Docket execution. - - Args: - docket: Docket instance - task_scope: Authorization scope - client_task_id: Client-provided task ID - - Returns: - Tuple of (execution, created_at, poll_interval_ms) - - Raises: - MCPError: If task not found or execution not found - """ - prefix = task_redis_prefix(task_scope) - task_meta_key = docket.key(f"{prefix}:{client_task_id}") - created_at_key = docket.key(f"{prefix}:{client_task_id}:created_at") - poll_interval_key = docket.key(f"{prefix}:{client_task_id}:poll_interval") - - # Fetch metadata (single round-trip with mget) - async with docket.redis() as redis: - task_key_bytes, created_at_bytes, poll_interval_bytes = await redis.mget( - task_meta_key, created_at_key, poll_interval_key - ) - - # Decode and validate task_key - task_key = task_key_bytes.decode("utf-8") if task_key_bytes else None - if not task_key: - raise MCPError(code=INVALID_PARAMS, message=f"Task {client_task_id} not found") - - # Get execution - execution = await docket.get_execution(task_key) - if not execution: - raise MCPError( - code=INVALID_PARAMS, - message=f"Task {client_task_id} execution not found", - ) - - # Parse metadata with defaults - created_at = created_at_bytes.decode("utf-8") if created_at_bytes else None - try: - poll_interval_ms = ( - int(poll_interval_bytes.decode("utf-8")) - if poll_interval_bytes - else DEFAULT_POLL_INTERVAL_MS - ) - except (ValueError, UnicodeDecodeError): - poll_interval_ms = DEFAULT_POLL_INTERVAL_MS - - return execution, created_at, poll_interval_ms - - -async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskResult: - """Handle MCP 'tasks/get' request (SEP-1686). - - Args: - server: FastMCP server instance - params: Request params containing taskId - - Returns: - GetTaskResult: Task status response with spec-compliant fields - """ - async with fastmcp.server.context.Context(fastmcp=server): - client_task_id = params.get("taskId") - if not client_task_id: - raise MCPError( - code=INVALID_PARAMS, message="Missing required parameter: taskId" - ) - - # Get authorization scope for task lookup - task_scope = get_task_scope() - - # Get Docket instance - docket = server._docket - if docket is None: - raise MCPError( - code=INTERNAL_ERROR, - message="Background tasks require Docket", - ) - - # Look up task execution and metadata - execution, created_at, poll_interval_ms = await _lookup_task_execution( - docket, task_scope, client_task_id - ) - - # Sync state from Redis - await execution.sync() - - # Map Docket state to MCP state - state_map = DOCKET_TO_MCP_STATE - mcp_state: Literal[ - "working", "input_required", "completed", "failed", "cancelled" - ] = state_map.get(execution.state, "failed") # type: ignore[assignment] # ty:ignore[invalid-assignment] - - # Build response (use default ttl since we don't track per-task values) - # createdAt is REQUIRED per SEP-1686 final spec (line 430) - # Per spec lines 447-448: SHOULD NOT include related-task metadata in tasks/get - error_message = None - status_message = None - - if execution.state == ExecutionState.FAILED: - try: - await execution.get_result(timeout=timedelta(seconds=0)) - except Exception as error: - error_message = str(error) - status_message = f"Task failed: {error_message}" - elif execution.progress and execution.progress.message: - # Extract progress message from Docket if available (spec line 403) - status_message = execution.progress.message - - # createdAt is required per spec, but can be None from Redis. The v2 - # Task model types createdAt/lastUpdatedAt as ISO 8601 strings, so - # normalize the stored value (or fall back to now) to an ISO string. - created_at_iso = _normalize_iso_timestamp(created_at) - - return GetTaskResult( - task_id=client_task_id, - status=mcp_state, - created_at=created_at_iso, - last_updated_at=datetime.now(timezone.utc).isoformat(), - ttl=DEFAULT_TTL_MS, - poll_interval=poll_interval_ms, - status_message=status_message, - ) - - -async def tasks_result_handler(server: FastMCP, params: dict[str, Any]) -> Any: - """Handle MCP 'tasks/result' request (SEP-1686). - - Converts raw task return values to MCP types based on task type. - - Args: - server: FastMCP server instance - params: Request params containing taskId - - Returns: - MCP result (CallToolResult, GetPromptResult, or ReadResourceResult) - """ - async with fastmcp.server.context.Context(fastmcp=server): - client_task_id = params.get("taskId") - if not client_task_id: - raise MCPError( - code=INVALID_PARAMS, message="Missing required parameter: taskId" - ) - - # Get authorization scope for task lookup - task_scope = get_task_scope() - - # Get execution from Docket (use instance attribute for cross-task access) - docket = server._docket - if docket is None: - raise MCPError( - code=INTERNAL_ERROR, - message="Background tasks require Docket", - ) - - # Look up full task key from Redis - task_meta_key = docket.key(f"{task_redis_prefix(task_scope)}:{client_task_id}") - async with docket.redis() as redis: - task_key_bytes = await redis.get(task_meta_key) - - task_key = None if task_key_bytes is None else task_key_bytes.decode("utf-8") - - if task_key is None: - raise MCPError( - code=INVALID_PARAMS, - message=f"Invalid taskId: {client_task_id} not found", - ) - - execution = await docket.get_execution(task_key) - if execution is None: - raise MCPError( - code=INVALID_PARAMS, - message=f"Invalid taskId: {client_task_id} not found", - ) - - # Sync state from Redis - await execution.sync() - - # Check if completed - state_map = DOCKET_TO_MCP_STATE - if execution.state not in (ExecutionState.COMPLETED, ExecutionState.FAILED): - mcp_state = state_map.get(execution.state, "failed") - raise MCPError( - code=INVALID_PARAMS, - message=f"Task not completed yet (current state: {mcp_state})", - ) - - # Get result from Docket - try: - raw_value = await execution.get_result(timeout=timedelta(seconds=0)) - except Exception as error: - # Task failed - return error result - return mcp_types.CallToolResult( - content=[mcp_types.TextContent(type="text", text=str(error))], - is_error=True, - _meta={ # type: ignore[call-arg] # _meta is Pydantic alias for meta field - "io.modelcontextprotocol/related-task": { - "taskId": client_task_id, - } - }, - ) - - # Parse task key to get component key - key_parts = parse_task_key(task_key) - component_key = key_parts["component_identifier"] - - # Look up component by its prefixed key (inlined from deleted get_component) - component: Tool | Resource | ResourceTemplate | Prompt | None = None - try: - if component_key.startswith("tool:"): - name, version_str = _parse_key_version(component_key[5:]) - version = VersionSpec(eq=version_str) if version_str else None - component = await server.get_tool(name, version) - elif component_key.startswith("resource:"): - uri, version_str = _parse_key_version(component_key[9:]) - version = VersionSpec(eq=version_str) if version_str else None - component = await server.get_resource(uri, version) - elif component_key.startswith("template:"): - uri, version_str = _parse_key_version(component_key[9:]) - version = VersionSpec(eq=version_str) if version_str else None - component = await server.get_resource_template(uri, version) - elif component_key.startswith("prompt:"): - name, version_str = _parse_key_version(component_key[7:]) - version = VersionSpec(eq=version_str) if version_str else None - component = await server.get_prompt(name, version) - except NotFoundError: - component = None - - if component is None: - raise MCPError( - code=INTERNAL_ERROR, - message=f"Component not found for task: {component_key}", - ) - - # Build related-task metadata - related_task_meta = { - "io.modelcontextprotocol/related-task": { - "taskId": client_task_id, - } - } - - # Convert based on component type. - # Each branch merges related_task_meta with any existing _meta - # (e.g. fastmcp.wrap_result) rather than overwriting it. - if isinstance(component, Tool): - if isinstance( - raw_value, mcp_types.InputRequiredResult | InputRequiredToolResult - ): - raise MCPError( - code=INTERNAL_ERROR, - message=( - f"Tool {component_key!r} requested input while running as a " - "background task. Input-required (multi-round-trip) tools " - "need a live request to answer the prompt and cannot run as " - "tasks; remove task execution from this tool or the code path " - "that returns an InputRequiredResult." - ), - ) - fastmcp_result = component.convert_result(raw_value) - mcp_result = fastmcp_result.to_mcp_result() - if isinstance(mcp_result, mcp_types.CallToolResult): - merged = {**(mcp_result.meta or {}), **related_task_meta} - mcp_result._meta = merged # type: ignore[attr-defined] # ty:ignore[unresolved-attribute] - elif isinstance(mcp_result, tuple): - content, structured_content = mcp_result - mcp_result = mcp_types.CallToolResult( - content=content, - structured_content=structured_content, - _meta=related_task_meta, # type: ignore[call-arg] # _meta is Pydantic alias for meta field - ) - else: - mcp_result = mcp_types.CallToolResult( - content=mcp_result, - _meta=related_task_meta, # type: ignore[call-arg] # _meta is Pydantic alias for meta field - ) - return mcp_result - - elif isinstance(component, Prompt): - fastmcp_result = component.convert_result(raw_value) - mcp_result = fastmcp_result.to_mcp_prompt_result() - merged = {**(mcp_result.meta or {}), **related_task_meta} - mcp_result._meta = merged # type: ignore[attr-defined] # ty:ignore[unresolved-attribute] - return mcp_result - - elif isinstance(component, ResourceTemplate): - fastmcp_result = component.convert_result(raw_value) - mcp_result = fastmcp_result.to_mcp_result(component.uri_template) - merged = {**(mcp_result.meta or {}), **related_task_meta} - mcp_result._meta = merged # type: ignore[attr-defined] # ty:ignore[unresolved-attribute] - return mcp_result - - elif isinstance(component, Resource): - fastmcp_result = component.convert_result(raw_value) - mcp_result = fastmcp_result.to_mcp_result(str(component.uri)) - merged = {**(mcp_result.meta or {}), **related_task_meta} - mcp_result._meta = merged # type: ignore[attr-defined] # ty:ignore[unresolved-attribute] - return mcp_result - - else: - raise MCPError( - code=INTERNAL_ERROR, - message=f"Internal error: Unknown component type: {type(component).__name__}", - ) - - -async def tasks_list_handler( - server: FastMCP, params: dict[str, Any] -) -> ListTasksResult: - """Handle MCP 'tasks/list' request (SEP-1686). - - Note: With client-side tracking, this returns minimal info. - - Args: - server: FastMCP server instance - params: Request params (cursor, limit) - - Returns: - ListTasksResult: Response with tasks list and pagination - """ - # Return empty list - client tracks tasks locally - return ListTasksResult(tasks=[], next_cursor=None) - - -async def tasks_cancel_handler( - server: FastMCP, params: dict[str, Any] -) -> CancelTaskResult: - """Handle MCP 'tasks/cancel' request (SEP-1686). - - Cancels a running task, transitioning it to cancelled state. - - Args: - server: FastMCP server instance - params: Request params containing taskId - - Returns: - CancelTaskResult: Task status response showing cancelled state - """ - async with fastmcp.server.context.Context(fastmcp=server): - client_task_id = params.get("taskId") - if not client_task_id: - raise MCPError( - code=INVALID_PARAMS, message="Missing required parameter: taskId" - ) - - # Get authorization scope for task lookup - task_scope = get_task_scope() - - # Get Docket instance - docket = server._docket - if docket is None: - raise MCPError( - code=INTERNAL_ERROR, - message="Background tasks require Docket", - ) - - # Look up task execution and metadata - execution, created_at, poll_interval_ms = await _lookup_task_execution( - docket, task_scope, client_task_id - ) - - # Cancel via Docket (now sets CANCELLED state natively) - # Note: We need to get task_key from execution.key for cancellation - await docket.cancel(execution.key) - - # Return task status with cancelled state - # createdAt is REQUIRED per SEP-1686 final spec (line 430) - # Per spec lines 447-448: SHOULD NOT include related-task metadata in tasks/cancel - return CancelTaskResult( - task_id=client_task_id, - status="cancelled", - created_at=_normalize_iso_timestamp(created_at), - last_updated_at=datetime.now(timezone.utc).isoformat(), - ttl=DEFAULT_TTL_MS, - poll_interval=poll_interval_ms, - status_message="Task cancelled", - ) diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/routing.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/routing.py deleted file mode 100644 index b8d7f0f4d..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/routing.py +++ /dev/null @@ -1,72 +0,0 @@ -"""Task routing helper for MCP components. - -Provides unified task mode enforcement and docket routing logic. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Literal - -import mcp_types -from mcp.shared.exceptions import MCPError -from mcp_types import METHOD_NOT_FOUND - -from fastmcp.utilities.tasks import TaskMeta -from fastmcp_tasks._legacy_wire.handlers import submit_to_docket - -if TYPE_CHECKING: - from fastmcp.prompts.base import Prompt - from fastmcp.resources.base import Resource - from fastmcp.resources.template import ResourceTemplate - from fastmcp.tools.base import Tool - -TaskType = Literal["tool", "resource", "template", "prompt"] - - -async def check_background_task( - component: Tool | Resource | ResourceTemplate | Prompt, - task_type: TaskType, - arguments: dict[str, Any] | None = None, - task_meta: TaskMeta | None = None, -) -> mcp_types.CreateTaskResult | None: - """Check task mode and submit to background if requested. - - Args: - component: The MCP component - task_type: Type of task ("tool", "resource", "template", "prompt") - arguments: Arguments for tool/prompt/template execution - task_meta: Task execution metadata. If provided, execute as background task. - - Returns: - CreateTaskResult if submitted to docket, None for sync execution - - Raises: - MCPError: If mode="required" but no task metadata, or mode="forbidden" - but task metadata is present - """ - task_config = component.task_config - - # Infer label from component - entity_label = f"{type(component).__name__} '{component.title or component.key}'" - - # Enforce mode="required" - must have task metadata - if task_config.mode == "required" and not task_meta: - raise MCPError( - code=METHOD_NOT_FOUND, - message=f"{entity_label} requires task-augmented execution", - ) - - # Enforce mode="forbidden" - cannot be called with task metadata - if not task_config.supports_tasks() and task_meta: - raise MCPError( - code=METHOD_NOT_FOUND, - message=f"{entity_label} does not support task-augmented execution", - ) - - # No task metadata - synchronous execution - if not task_meta: - return None - - # fn_key is expected to be set; fall back to component.key for direct calls - fn_key = task_meta.fn_key or component.key - return await submit_to_docket(task_type, fn_key, component, arguments, task_meta) diff --git a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/subscriptions.py b/fastmcp_tasks/fastmcp_tasks/_legacy_wire/subscriptions.py deleted file mode 100644 index 05526a6f4..000000000 --- a/fastmcp_tasks/fastmcp_tasks/_legacy_wire/subscriptions.py +++ /dev/null @@ -1,282 +0,0 @@ -"""Task subscription helpers for sending MCP notifications (SEP-1686). - -Subscribes to Docket execution state changes and sends notifications/tasks/status -to clients when their tasks change state. - -This module requires fastmcp[tasks] (pydocket). It is only imported when docket is available. -""" - -from __future__ import annotations - -import asyncio -from contextlib import suppress -from datetime import datetime, timezone -from typing import TYPE_CHECKING - -from docket.execution import ExecutionState -from mcp_types import TaskStatusNotification, TaskStatusNotificationParams - -from fastmcp.utilities.logging import get_logger -from fastmcp.utilities.tasks import DEFAULT_TTL_MS -from fastmcp_tasks._legacy_wire.requests import DOCKET_TO_MCP_STATE -from fastmcp_tasks.keys import parse_task_key, task_redis_prefix - -if TYPE_CHECKING: - from docket import Docket - from docket.execution import Execution - from mcp.server.session import ServerSession - -logger = get_logger(__name__) - -# Initial interval for reconciling execution state against Redis (seconds). The -# interval doubles on each idle reconcile up to the task's poll_interval, so fast -# tasks are caught within the first ~20ms checks while long-running tasks converge -# to roughly one sync per advertised poll interval. -_MIN_RECONCILE_INTERVAL_SECONDS = 0.02 - - -async def subscribe_to_task_updates( - task_id: str, - task_key: str, - session: ServerSession, - docket: Docket, - poll_interval_ms: int = 5000, -) -> None: - """Subscribe to Docket execution events and send MCP notifications. - - Per SEP-1686 lines 436-444, servers MAY send notifications/tasks/status - when task state changes. This is an optional optimization that reduces - client polling frequency. - - Args: - task_id: Client-visible task ID (server-generated UUID) - task_key: Internal Docket execution key (includes session, type, component) - session: MCP ServerSession for sending notifications - docket: Docket instance for subscribing to execution events - poll_interval_ms: Poll interval in milliseconds to include in notifications - - Note: Docket's ``execution.subscribe()`` replays the current state and a progress - event before it subscribes to Redis pub/sub. A task that completes during that - window has its terminal state publish lost, so no live event ever arrives — a - common case for fast tasks. Because there is no reliable signal for when the - subscription goes live (the replayed state event arrives two iterations early), - we simply reconcile the execution against Redis on every idle interval until a - terminal state is observed. The interval backs off exponentially toward the - task's advertised poll interval, so a long-running task costs about one sync per - poll interval while live pub/sub events still short-circuit the wait instantly. - """ - terminal_states = { - ExecutionState.COMPLETED, - ExecutionState.FAILED, - ExecutionState.CANCELLED, - } - try: - execution = await docket.get_execution(task_key) - if execution is None: - logger.warning(f"No execution found for task {task_id}") - return - - subscription = execution.subscribe() - # Keep a single outstanding __anext__ across reconcile timeouts. asyncio.wait - # returns on timeout without cancelling it, so the generator (and its pub/sub - # subscription) stays intact — unlike wait_for, which would cancel mid-iteration. - next_event = asyncio.ensure_future(subscription.__anext__()) - # Reconcile cadence backs off exponentially so a task that runs for a long - # time (or that no worker ever claims) doesn't pin this loop at 50 syncs/sec - # forever; the task's advertised poll interval is the natural ceiling. - reconcile_backoff = _MIN_RECONCILE_INTERVAL_SECONDS - reconcile_ceiling = max( - poll_interval_ms / 1000, _MIN_RECONCILE_INTERVAL_SECONDS - ) - try: - while True: - done, _ = await asyncio.wait({next_event}, timeout=reconcile_backoff) - if not done: - # No live event yet: reconcile against Redis in case a - # terminal transition was published before pub/sub went live. - await execution.sync() - if execution.state in terminal_states: - await _send_status_notification( - session=session, - task_id=task_id, - task_key=task_key, - docket=docket, - state=execution.state, - poll_interval_ms=poll_interval_ms, - ) - break - reconcile_backoff = min(reconcile_backoff * 2, reconcile_ceiling) - continue - - try: - event = next_event.result() - except StopAsyncIteration: - break - - if event["type"] == "state": - state = ExecutionState(event["state"]) - # Send notifications/tasks/status when state changes - await _send_status_notification( - session=session, - task_id=task_id, - task_key=task_key, - docket=docket, - state=state, - poll_interval_ms=poll_interval_ms, - ) - # Stop subscribing once the task reaches a terminal state - if state in terminal_states: - break - elif event["type"] == "progress": - # Send notification when progress message changes - await _send_progress_notification( - session=session, - task_id=task_id, - task_key=task_key, - docket=docket, - execution=execution, - poll_interval_ms=poll_interval_ms, - ) - - next_event = asyncio.ensure_future(subscription.__anext__()) - finally: - if not next_event.done(): - next_event.cancel() - with suppress(asyncio.CancelledError, StopAsyncIteration): - await next_event - await subscription.aclose() - - except Exception as e: - logger.warning(f"Subscription task failed for {task_id}: {e}", exc_info=True) - - -async def _send_status_notification( - session: ServerSession, - task_id: str, - task_key: str, - docket: Docket, - state: ExecutionState, - poll_interval_ms: int = 5000, -) -> None: - """Send notifications/tasks/status to client. - - Per SEP-1686 line 454: notification SHOULD NOT include related-task metadata - (taskId is already in params). - - Args: - session: MCP ServerSession - task_id: Client-visible task ID - task_key: Internal task key (for metadata lookup) - docket: Docket instance - state: Docket execution state (enum) - poll_interval_ms: Poll interval in milliseconds - """ - # Map Docket state to MCP status - state_map = DOCKET_TO_MCP_STATE - mcp_status = state_map.get(state, "failed") - - # Extract task_scope from task_key for Redis lookup - key_parts = parse_task_key(task_key) - task_scope = key_parts["task_scope"] - - created_at_key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:created_at") - async with docket.redis() as redis: - created_at_bytes = await redis.get(created_at_key) - - created_at = ( - created_at_bytes.decode("utf-8") - if created_at_bytes - else datetime.now(timezone.utc).isoformat() - ) - - # Build status message - status_message = None - if state == ExecutionState.COMPLETED: - status_message = "Task completed successfully" - elif state == ExecutionState.FAILED: - status_message = "Task failed" - elif state == ExecutionState.CANCELLED: - status_message = "Task cancelled" - - params_dict = { - "taskId": task_id, - "status": mcp_status, - "createdAt": created_at, - "lastUpdatedAt": datetime.now(timezone.utc).isoformat(), - "ttl": DEFAULT_TTL_MS, - "pollInterval": poll_interval_ms, - } - - if status_message: - params_dict["statusMessage"] = status_message - - # Create notification (no related-task metadata per spec line 454) - notification = TaskStatusNotification( - params=TaskStatusNotificationParams.model_validate(params_dict), - ) - - # Send notification (don't let failures break the subscription) - with suppress(Exception): - await session.send_notification(notification) # type: ignore[arg-type] # ty:ignore[invalid-argument-type] - - -async def _send_progress_notification( - session: ServerSession, - task_id: str, - task_key: str, - docket: Docket, - execution: Execution, - poll_interval_ms: int = 5000, -) -> None: - """Send notifications/tasks/status when progress updates. - - Args: - session: MCP ServerSession - task_id: Client-visible task ID - task_key: Internal task key - docket: Docket instance - execution: Execution object with current progress - poll_interval_ms: Poll interval in milliseconds - """ - # Sync execution to get latest progress - await execution.sync() - - # Only send if there's a progress message - if not execution.progress or not execution.progress.message: - return - - # Map Docket state to MCP status - state_map = DOCKET_TO_MCP_STATE - mcp_status = state_map.get(execution.state, "failed") - - # Extract task_scope from task_key for Redis lookup - key_parts = parse_task_key(task_key) - task_scope = key_parts["task_scope"] - - created_at_key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:created_at") - async with docket.redis() as redis: - created_at_bytes = await redis.get(created_at_key) - - created_at = ( - created_at_bytes.decode("utf-8") - if created_at_bytes - else datetime.now(timezone.utc).isoformat() - ) - - params_dict = { - "taskId": task_id, - "status": mcp_status, - "createdAt": created_at, - "lastUpdatedAt": datetime.now(timezone.utc).isoformat(), - "ttl": DEFAULT_TTL_MS, - "pollInterval": poll_interval_ms, - "statusMessage": execution.progress.message, - } - - # Create and send notification - notification = TaskStatusNotification( - params=TaskStatusNotificationParams.model_validate(params_dict), - ) - - with suppress(Exception): - await session.send_notification(notification) # type: ignore[arg-type] # ty:ignore[invalid-argument-type] diff --git a/fastmcp_tasks/fastmcp_tasks/components.py b/fastmcp_tasks/fastmcp_tasks/components.py index 537e6e645..0f3c0370a 100644 --- a/fastmcp_tasks/fastmcp_tasks/components.py +++ b/fastmcp_tasks/fastmcp_tasks/components.py @@ -3,8 +3,8 @@ During the SEP-1686 -> SEP-2663 migration the ``register_with_docket`` / ``add_to_docket`` / ``coerce_task_arguments`` methods were removed from the core ``FastMCPComponent`` classes (Tool, Resource, ResourceTemplate, Prompt). Their -bodies are preserved here verbatim as type-dispatched functions so Phase 3 can -wire them into ``TasksExtension`` without reconstructing the calling conventions. +bodies live here as type-dispatched functions that ``TasksExtension`` wires into +the Docket engine, preserving each type's calling convention. The functions dispatch on the concrete component type because each type splats its arguments differently into the Docket-registered callable: @@ -15,9 +15,9 @@ its arguments differently into the Docket-registered callable: - Base ``Tool``/``Resource``/``ResourceTemplate``/``Prompt`` register their ``run``/``read``/``render`` entry point and pass arguments positionally. -Only tools carry a task-capable ``task_config`` after the migration (SEP-2663 is -tools-only); the resource/prompt/template branches are retained for engine -completeness and Phase 3's decision, not because core still declares them. +Only tools carry a task-capable ``task_config`` (SEP-2663 is tools-only); the +resource/prompt/template branches are retained for engine completeness, not +because core still declares them. """ from __future__ import annotations diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index c18e33309..ff56097c7 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/fastmcp_tasks/fastmcp_tasks/context.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from docket import Docket from mcp.server.session import ServerSession + from fastmcp.server.context import Context from fastmcp.server.server import FastMCP _logger = logging.getLogger(__name__) @@ -362,3 +363,48 @@ def get_task_server(task_id: str) -> FastMCP | None: if server is None: _task_server_map.pop(task_id, None) return server + + +def resolve_worker_server() -> FastMCP | None: + """Return the server owning the current task's tool, or None outside a task. + + Installed as core's worker-server resolver by ``TasksExtension`` so + ``get_server()``/``CurrentFastMCP()`` inside a worker resolve to the (child) + server the task was submitted against, not the root that runs the worker. + """ + task_info = get_task_context() + if task_info is None: + return None + return get_task_server(task_info.task_id) + + +async def make_task_context() -> Context | None: + """Build and enter a worker ``Context`` for the current background task. + + Installed as core's background-context factory by ``TasksExtension`` so a + ``ctx: Context`` parameter resolves inside a Docket worker. Returns ``None`` + when not running in a task (so core falls through to its usual error). The + snapshot restored by ``restore_task_snapshot`` supplies the origin request + id; the server prefers the one registered at submission time so mounted + tasks resolve to the child server. No live session is attached — SEP-2663 + input and status are polled, so the worker needs no back-channel. + """ + from fastmcp.server.context import Context + from fastmcp.server.dependencies import get_server + + task_info = get_task_context() + if task_info is None: + return None + + server = get_task_server(task_info.task_id) or get_server() + snapshot = _recall_snapshot(task_info.task_id) + origin_request_id = snapshot.origin_request_id if snapshot else None + + ctx = Context( + fastmcp=server, + session=None, + task_id=task_info.task_id, + origin_request_id=origin_request_id, + ) + await ctx.__aenter__() + return ctx diff --git a/fastmcp_tasks/fastmcp_tasks/creation.py b/fastmcp_tasks/fastmcp_tasks/creation.py new file mode 100644 index 000000000..9409c9117 --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/creation.py @@ -0,0 +1,193 @@ +"""SEP-2663 task creation: enqueue an augmented tool call to Docket. + +Adapted from the SEP-1686 ``submit_to_docket`` path. The wire surface changed +(a flat ``CreateTaskResult`` with ``ttlMs``/``pollIntervalMs``, no client-supplied +task id or ttl) and the SEP-1686 push machinery — the initial status +notification, the per-task subscription, and the notification subscriber — is +gone, because SEP-2663 in-task input and status are polled, not pushed. The +operational core is preserved: strict argument coercion up front, a +server-generated high-entropy task id, the auth-scoped compound key, the context +snapshot restored in the worker, and durable creation (metadata is written +before the result is returned, so a subsequent ``tasks/get`` always resolves). +""" + +from __future__ import annotations + +import asyncio +import secrets +from datetime import datetime, timezone +from typing import TYPE_CHECKING + +from mcp.shared.exceptions import MCPError +from mcp_types import INTERNAL_ERROR + +from fastmcp.tools.base import Tool +from fastmcp.tools.function_tool import _strict_input_validation +from fastmcp.utilities.logging import get_logger +from fastmcp_tasks.components import add_component_to_docket, coerce_task_arguments +from fastmcp_tasks.context import ( + TaskContextSnapshot, + get_task_scope, + register_task_server, +) +from fastmcp_tasks.dependencies import _current_docket +from fastmcp_tasks.keys import build_task_key, task_redis_prefix +from fastmcp_tasks.models import CreateTaskResult + +if TYPE_CHECKING: + from docket import Docket + + from fastmcp.server.context import Context + from fastmcp.server.server import FastMCP + +logger = get_logger(__name__) + +# Redis mapping TTL buffer: keep task metadata a little longer than the Docket +# execution TTL so a client polling right at the edge still resolves the task. +TASK_MAPPING_TTL_BUFFER_SECONDS = 15 * 60 + +# Bounded read-your-writes wait so durable creation holds on distributed +# backends where the enqueued execution may not be immediately visible. +_DURABLE_CREATE_TIMEOUT_SECONDS = 5.0 +_DURABLE_CREATE_POLL_SECONDS = 0.02 + + +async def create_task( + tool: Tool, + arguments: dict[str, object] | None, + context: Context, +) -> CreateTaskResult: + """Run an augmented ``tools/call`` as a background task (SEP-2663). + + Coerces and validates arguments (honoring strict input validation), mints a + server-generated task id, snapshots the request context, enqueues the tool's + callable on Docket under the auth-scoped compound key, and returns a + ``CreateTaskResult`` in ``working`` status. Does not return until the task's + metadata is durably written and its execution is visible, so an immediately + following ``tasks/get`` resolves. + """ + # The interceptor resolves the tool via get_tool(), which for a mounted tool + # returns a provider wrapper — but Docket registered the underlying component + # from get_tasks() under the same key, with that component's calling + # convention (a FunctionTool splats **kwargs; a base Tool takes the dict + # positionally). Execute against the registered component so coercion and + # argument-splatting match what the worker will invoke. + component = await _registered_task_component(context, tool) + + coerced = coerce_task_arguments( + component, dict(arguments or {}), strict=_strict_input_validation() + ) + + task_id = secrets.token_urlsafe(32) + created_at = datetime.now(timezone.utc).isoformat() + + task_scope = get_task_scope() + + docket = context.fastmcp._docket or _current_docket.get() + if docket is None: + raise MCPError( + code=INTERNAL_ERROR, + message="Background tasks require a running tasks extension (Docket).", + ) + + # Resolve mounted tasks to the owning (child) server in the worker, so + # CurrentFastMCP()/ctx.fastmcp inside the task point at the server the tool + # lives on rather than the root the interceptor ran on (#3571). + register_task_server(task_id, _owning_server(tool, context.fastmcp)) + + key = component.key + task_key = build_task_key(task_scope, task_id, "tool", key) + + ttl_ms = int(docket.execution_ttl.total_seconds() * 1000) + ttl_seconds = int(ttl_ms / 1000) + TASK_MAPPING_TTL_BUFFER_SECONDS + poll_interval_ms = int(component.task_config.poll_interval.total_seconds() * 1000) + + prefix = task_redis_prefix(task_scope) + task_meta_key = docket.key(f"{prefix}:{task_id}") + created_at_key = docket.key(f"{prefix}:{task_id}:created_at") + poll_interval_key = docket.key(f"{prefix}:{task_id}:poll_interval") + + snapshot = TaskContextSnapshot.capture() + + async with docket.redis() as redis: + await redis.set(task_meta_key, task_key, ex=ttl_seconds) + await redis.set(created_at_key, created_at, ex=ttl_seconds) + await redis.set(poll_interval_key, str(poll_interval_ms), ex=ttl_seconds) + + await snapshot.save(docket, task_scope, task_id, ttl_seconds) + + await add_component_to_docket( + component, docket, coerced, fn_key=key, task_key=task_key + ) + + await _await_durable_creation(docket, task_key) + + return CreateTaskResult( + task_id=task_id, + status="working", + created_at=created_at, + last_updated_at=created_at, + ttl_ms=ttl_ms, + poll_interval_ms=poll_interval_ms, + ) + + +def _owning_server(tool: Tool, fallback: FastMCP) -> FastMCP: + """The server a mounted tool lives on, for worker context resolution. + + A mounted tool is a ``FastMCPProviderTool`` that references the child server + it came from, so ``CurrentFastMCP()``/``ctx.fastmcp`` inside the task point at + that server rather than the root the interceptor ran on (#3571). Resolution + is single-level: a tool reached through several nested mounts resolves to the + outermost mounted child (the mount point the call arrived through), which + still reaches deeper components through its own mounts. A non-mounted tool + falls back to the server the call arrived on. + """ + from fastmcp.server.providers.fastmcp_provider import FastMCPProviderTool + + if isinstance(tool, FastMCPProviderTool): + return tool._server + return fallback + + +async def _registered_task_component(context: Context, tool: Tool) -> Tool: + """Return the component Docket registered for ``tool``'s key. + + ``get_tasks()`` yields the same components that were registered with Docket + (the underlying ``FunctionTool`` for a mounted tool, not the provider + wrapper the interceptor's ``get_tool`` returns). Matching by ``key`` recovers + the registered component so the calling convention agrees with the worker. + Falls back to the interceptor's tool if no match is found (e.g. a dynamically + added tool not present at registration time). + """ + for component in await context.fastmcp.get_tasks(): + if component.key == tool.key and isinstance(component, Tool): + return component + return tool + + +async def _await_durable_creation(docket: Docket, task_key: str) -> None: + """Block until the enqueued execution is visible (durable-create MUST). + + The metadata write above already makes ``tasks/get`` resolvable; this extra + check guards distributed backends where the execution record propagates + slightly behind the enqueue. Bounded so a backend hiccup can't hang creation. + """ + deadline = asyncio.get_event_loop().time() + _DURABLE_CREATE_TIMEOUT_SECONDS + while True: + execution = await docket.get_execution(task_key) + if execution is not None: + return + if asyncio.get_event_loop().time() >= deadline: + # SEP-2663 durable-create: a CreateTaskResult MUST NOT be returned + # unless a subsequent tasks/get would resolve. Returning a handle + # that can 404 is the exact failure the requirement forbids, so a + # backend that never surfaces the execution is a create error. + raise MCPError( + code=INTERNAL_ERROR, + message=( + "Task creation did not become durable in time; the task " + "backend did not surface the enqueued execution." + ), + ) + await asyncio.sleep(_DURABLE_CREATE_POLL_SECONDS) diff --git a/fastmcp_tasks/fastmcp_tasks/dependencies.py b/fastmcp_tasks/fastmcp_tasks/dependencies.py index bb082483d..0598fe6af 100644 --- a/fastmcp_tasks/fastmcp_tasks/dependencies.py +++ b/fastmcp_tasks/fastmcp_tasks/dependencies.py @@ -4,7 +4,7 @@ Moved out of ``fastmcp.server.dependencies`` during the SEP-1686 -> SEP-2663 migration. These helpers are all docket-touching: the ``require_docket`` install-hint, the docket/worker ContextVars, and the ``CurrentDocket`` / ``CurrentWorker`` dependencies. Everything here is wire-agnostic engine plumbing -that Phase 3 rewires into ``TasksExtension``. +that ``TasksExtension`` drives. The generic ``is_docket_available`` probe stays in ``fastmcp.server.dependencies`` (core's ``Context``/``Progress`` still use it) and is re-exported here for the diff --git a/fastmcp_tasks/fastmcp_tasks/extension.py b/fastmcp_tasks/fastmcp_tasks/extension.py new file mode 100644 index 000000000..646baca3c --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/extension.py @@ -0,0 +1,262 @@ +"""The SEP-2663 tasks extension: `io.modelcontextprotocol/tasks`. + +`TasksExtension` is the wire adapter that turns FastMCP's task engine into an +`io.modelcontextprotocol/tasks` server extension. Registering it enables +`task=True` tools: + +```python +from fastmcp import FastMCP +from fastmcp_tasks import TasksExtension + +mcp = FastMCP("Server") +mcp.add_extension(TasksExtension(url="redis://localhost:6379/0")) + + +@mcp.tool(task=True) +async def crunch(dataset: str) -> str: + ... +``` + +The extension contributes the negotiated capability, the three additive request +methods (`tasks/get`, `tasks/update`, `tasks/cancel`), a `tools/call` interceptor +that decides whether to run a call as a task, and a lifespan that starts the +Docket backend/worker and installs the worker-side `Context` hooks core exposes. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Sequence +from contextlib import asynccontextmanager +from datetime import timedelta +from typing import TYPE_CHECKING, Any + +from mcp.server.context import ServerRequestContext +from mcp.shared.exceptions import MCPError +from mcp_types.version import MODERN_PROTOCOL_VERSIONS + +from fastmcp.exceptions import NotFoundError +from fastmcp.server.extensions import MethodBinding, ServerExtension +from fastmcp.utilities.logging import get_logger +from fastmcp.utilities.tasks import TASKS_EXTENSION_ID +from fastmcp_tasks.creation import create_task +from fastmcp_tasks.handlers import tasks_cancel, tasks_get, tasks_update +from fastmcp_tasks.models import ( + MISSING_REQUIRED_CLIENT_CAPABILITY, + CancelTaskParams, + CancelTaskResult, + GetTaskParams, + GetTaskResult, + UpdateTaskParams, + UpdateTaskResult, + missing_capability_error_data, +) +from fastmcp_tasks.settings import DocketSettings + +if TYPE_CHECKING: + import mcp_types + + from fastmcp.server.context import Context + from fastmcp.server.extensions import ToolCallContinuation, ToolCallOutcome + +logger = get_logger(__name__) + +# SEP-2663's request methods exist only at the 2026-07-28 era (the extensions +# mechanism itself is era-gated). Off that era the methods report as not found. +_TASK_METHOD_VERSIONS = frozenset(MODERN_PROTOCOL_VERSIONS) + + +class TasksExtension(ServerExtension): + """FastMCP server extension implementing SEP-2663 background tasks. + + Construct with backend/worker configuration; anything omitted falls back to + the ``FASTMCP_DOCKET_*`` environment defaults (unchanged from FastMCP 3), so + ``TasksExtension()`` works out of the box on an env-configured deployment. + """ + + identifier = TASKS_EXTENSION_ID + + def __init__( + self, + *, + url: str | None = None, + name: str | None = None, + worker_name: str | None = None, + concurrency: int | None = None, + redelivery_timeout: timedelta | None = None, + reconnection_delay: timedelta | None = None, + minimum_check_interval: timedelta | None = None, + ) -> None: + overrides: dict[str, Any] = { + "url": url, + "name": name, + "worker_name": worker_name, + "concurrency": concurrency, + "redelivery_timeout": redelivery_timeout, + "reconnection_delay": reconnection_delay, + "minimum_check_interval": minimum_check_interval, + } + self._settings = DocketSettings( + **{k: v for k, v in overrides.items() if v is not None} + ) + + @property + def docket_settings(self) -> DocketSettings: + """The resolved Docket settings (backend URL, worker options).""" + return self._settings + + def settings(self) -> dict[str, Any]: + """The tasks extension advertises no per-extension settings.""" + return {} + + def methods(self) -> Sequence[MethodBinding]: + return [ + MethodBinding( + method="tasks/get", + params_type=GetTaskParams, + handler=self._handle_get, + protocol_versions=_TASK_METHOD_VERSIONS, + ), + MethodBinding( + method="tasks/update", + params_type=UpdateTaskParams, + handler=self._handle_update, + protocol_versions=_TASK_METHOD_VERSIONS, + ), + MethodBinding( + method="tasks/cancel", + params_type=CancelTaskParams, + handler=self._handle_cancel, + protocol_versions=_TASK_METHOD_VERSIONS, + ), + ] + + async def _handle_get( + self, ctx: ServerRequestContext[Any, Any], params: GetTaskParams + ) -> GetTaskResult: + return await tasks_get(self.server, params.task_id) + + async def _handle_update( + self, ctx: ServerRequestContext[Any, Any], params: UpdateTaskParams + ) -> UpdateTaskResult: + return await tasks_update(self.server, params.task_id, params.input_responses) + + async def _handle_cancel( + self, ctx: ServerRequestContext[Any, Any], params: CancelTaskParams + ) -> CancelTaskResult: + return await tasks_cancel(self.server, params.task_id) + + async def intercept_tool_call( + self, + params: mcp_types.CallToolRequestParams, + context: Context, + call_next: ToolCallContinuation, + ) -> ToolCallOutcome: + """Decide whether to run this ``tools/call`` as a task. + + Consults the tool's ``TaskConfig`` mode and the client's per-request + opt-in: ``required`` always tasks (raising -32003 if the client did not + opt in), ``optional`` tasks only when the client opted in, ``forbidden`` + never tasks. A non-task call passes straight through to the tool body. + """ + try: + tool = await context.fastmcp.get_tool(params.name) + except NotFoundError: + tool = None + if tool is None or not tool.task_config.supports_tasks(): + return await call_next() + + # Extension negotiation exists only on the modern era: the SDK strips + # `capabilities.extensions` from pre-2026 handshakes, so a legacy client + # cannot have negotiated this extension — a `_meta` opt-in arriving on a + # handshake-era connection is treated as absent. This also keeps a + # `CreateTaskResult` off legacy connections, whose result validation + # does not admit it. + rc = context.request_context + on_modern_era = ( + rc is not None and rc.protocol_version in MODERN_PROTOCOL_VERSIONS + ) + opted_in = ( + on_modern_era + and context.client_extension_settings(TASKS_EXTENSION_ID) is not None + ) + mode = tool.task_config.mode + + if mode == "required": + if not opted_in: + raise MCPError( + code=MISSING_REQUIRED_CLIENT_CAPABILITY, + message=( + f"Tool {tool.name!r} requires the tasks extension " + f"({TASKS_EXTENSION_ID}); the client did not declare it " + "for this request." + ), + data=missing_capability_error_data(), + ) + return await create_task(tool, params.arguments, context) + + if mode == "optional" and opted_in: + return await create_task(tool, params.arguments, context) + + return await call_next() + + @asynccontextmanager + async def lifespan(self) -> AsyncIterator[None]: + """Start the Docket backend/worker and install the worker-side hooks. + + Installs core's background-context factory and in-task elicitation + handler for the duration so a worker's ``ctx`` (progress, elicitation) + functions, then runs the Docket lifespan. The hooks are process-global + and refcounted: with several servers in one process (each its own + runtime-tree root), the hooks stay installed until the last tasks + extension shuts down, so one server's exit cannot strand another + server's in-flight workers. + """ + from fastmcp_tasks.lifespan import docket_lifespan + + _install_worker_hooks() + try: + async with docket_lifespan(self.server, self._settings): + yield + finally: + _release_worker_hooks() + + +# The worker-side hooks core exposes are process-global, but several servers in +# one process may each run a TasksExtension (sibling roots in tests, or two +# apps sharing an interpreter). Refcount the installs so the hooks are cleared +# only when the last active extension lifespan exits. The installed callables +# are stateless module functions that resolve their target per task, so +# repeated installs are idempotent. +_active_worker_hook_holds: int = 0 + + +def _install_worker_hooks() -> None: + from fastmcp.server.context import set_task_elicitation_handler + from fastmcp.server.dependencies import ( + set_background_context_factory, + set_worker_server_resolver, + ) + from fastmcp_tasks.context import make_task_context, resolve_worker_server + from fastmcp_tasks.input_store import elicit_in_task + + global _active_worker_hook_holds + _active_worker_hook_holds += 1 + set_background_context_factory(make_task_context) + set_worker_server_resolver(resolve_worker_server) + set_task_elicitation_handler(elicit_in_task) + + +def _release_worker_hooks() -> None: + from fastmcp.server.context import set_task_elicitation_handler + from fastmcp.server.dependencies import ( + set_background_context_factory, + set_worker_server_resolver, + ) + + global _active_worker_hook_holds + _active_worker_hook_holds -= 1 + if _active_worker_hook_holds <= 0: + _active_worker_hook_holds = 0 + set_task_elicitation_handler(None) + set_worker_server_resolver(None) + set_background_context_factory(None) diff --git a/fastmcp_tasks/fastmcp_tasks/handlers.py b/fastmcp_tasks/fastmcp_tasks/handlers.py new file mode 100644 index 000000000..e7c906b8e --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/handlers.py @@ -0,0 +1,283 @@ +"""SEP-2663 task query/management handlers: tasks/get, tasks/update, tasks/cancel. + +Adapted from the SEP-1686 ``requests.py``. The three CRUD-ish handlers survive, +reshaped to the new wire: + +- ``tasks/get`` merges the old ``tasks/get`` and ``tasks/result``: the finished + result is *inlined* into the response for a completed task, a JSON-RPC-shaped + ``error`` for a failed one, and the outstanding ``inputRequests`` for a task + waiting on input. +- ``tasks/update`` is new: it delivers ``inputResponses`` to the in-task input + store, resuming a parked worker. +- ``tasks/cancel`` returns an empty ack (SEP-2663) instead of a task snapshot. +- ``tasks/list`` and ``tasks/result`` are gone (removed by SEP-2663). + +The auth-scoped compound key is the authorization boundary: a request resolves a +task only under its own scope's Redis prefix, so a scope mismatch is +indistinguishable from a missing task (both raise -32602 "Task not found"), +which avoids leaking task existence across callers. +""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Any, Literal + +import mcp_types +from docket.execution import ExecutionState +from mcp.shared.exceptions import MCPError +from mcp_types import INVALID_PARAMS + +from fastmcp.exceptions import NotFoundError +from fastmcp.tools.base import InputRequiredToolResult, Tool +from fastmcp.utilities.tasks import DEFAULT_POLL_INTERVAL_MS +from fastmcp.utilities.versions import VersionSpec +from fastmcp_tasks.context import get_task_scope +from fastmcp_tasks.input_store import deliver_input_responses, read_outstanding_inputs +from fastmcp_tasks.keys import parse_task_key, task_redis_prefix +from fastmcp_tasks.models import ( + CancelTaskResult, + GetTaskResult, + UpdateTaskResult, +) + +if TYPE_CHECKING: + from docket import Docket + + from fastmcp.server.server import FastMCP + +# Docket execution state -> SEP-2663 task status. `input_required` is not a +# Docket state; it is derived from the in-task input store (see tasks_get). +DOCKET_TO_MCP_STATE: dict[ExecutionState, str] = { + ExecutionState.SCHEDULED: "working", + ExecutionState.QUEUED: "working", + ExecutionState.RUNNING: "working", + ExecutionState.COMPLETED: "completed", + ExecutionState.FAILED: "failed", + ExecutionState.CANCELLED: "cancelled", +} + +_WORKING_STATES = frozenset( + {ExecutionState.SCHEDULED, ExecutionState.QUEUED, ExecutionState.RUNNING} +) + + +def _task_not_found(task_id: str) -> MCPError: + """The single "not found" error for missing, expired, or cross-scope ids. + + Uses one message for all three so a caller cannot probe another scope's task + ids by distinguishing "not yours" from "does not exist". + """ + return MCPError(code=INVALID_PARAMS, message=f"Task {task_id} not found") + + +def _normalize_iso_timestamp(stored: str | None) -> str: + """Return an ISO 8601 timestamp for createdAt, tolerating a missing value.""" + if stored: + try: + return datetime.fromisoformat(stored.replace("Z", "+00:00")).isoformat() + except (ValueError, AttributeError): + pass + return datetime.now(timezone.utc).isoformat() + + +def _parse_key_version(key_suffix: str) -> tuple[str, str | None]: + """Split a component key suffix into (name, version) on the last ``@``.""" + if "@" not in key_suffix: + return key_suffix, None + name, version = key_suffix.rsplit("@", 1) + return name, version if version else None + + +def _ttl_ms(docket: Docket) -> int: + """The task TTL in milliseconds, from Docket's execution TTL (server-set).""" + return int(docket.execution_ttl.total_seconds() * 1000) + + +async def _lookup_task( + docket: Docket, task_scope: str | None, task_id: str +) -> tuple[Any, str, str | None, int]: + """Resolve a task's execution and stored metadata within the caller's scope. + + Returns ``(execution, task_key, created_at, poll_interval_ms)``. Raises the + shared "not found" error when the scope-prefixed metadata is absent or the + execution has expired. + """ + prefix = task_redis_prefix(task_scope) + meta_key = docket.key(f"{prefix}:{task_id}") + created_at_key = docket.key(f"{prefix}:{task_id}:created_at") + poll_key = docket.key(f"{prefix}:{task_id}:poll_interval") + + async with docket.redis() as redis: + # Docket's Redis client mirrors redis-py's variadic ``mget(*keys)`` at + # runtime; its type stub declares a single ``Sequence`` arg, so the + # positional form is correct but needs a targeted ignore. + values = await redis.mget(meta_key, created_at_key, poll_key) # ty: ignore[too-many-positional-arguments] + task_key_bytes, created_at_bytes, poll_bytes = values + + task_key = task_key_bytes.decode("utf-8") if task_key_bytes else None + if not task_key: + raise _task_not_found(task_id) + + execution = await docket.get_execution(task_key) + if not execution: + raise _task_not_found(task_id) + + created_at = created_at_bytes.decode("utf-8") if created_at_bytes else None + + try: + poll_interval_ms = ( + int(poll_bytes.decode("utf-8")) if poll_bytes else DEFAULT_POLL_INTERVAL_MS + ) + except (ValueError, UnicodeDecodeError): + poll_interval_ms = DEFAULT_POLL_INTERVAL_MS + + return execution, task_key, created_at, poll_interval_ms + + +async def _resolve_tool(server: FastMCP, task_key: str) -> Tool: + """Resolve the Tool a task ran, from its compound key (tools-only surface).""" + component_key = parse_task_key(task_key)["component_identifier"] + if not component_key.startswith("tool:"): + raise MCPError( + code=mcp_types.INTERNAL_ERROR, + message=f"Task component is not a tool: {component_key}", + ) + name, version_str = _parse_key_version(component_key[len("tool:") :]) + version = VersionSpec(eq=version_str) if version_str else None + try: + tool = await server.get_tool(name, version) + except NotFoundError: + tool = None + if tool is None: + raise MCPError( + code=mcp_types.INTERNAL_ERROR, + message=f"Component not found for task: {component_key}", + ) + return tool + + +def _inline_result(tool: Tool, raw_value: Any) -> dict[str, Any]: + """Convert a completed task's raw return into an inlined CallToolResult dict. + + A guard tool that returned an ``InputRequiredResult`` from inside a task is + rejected: multi-round-trip guards need a live request to answer the prompt + and cannot complete as a task. + """ + if isinstance(raw_value, mcp_types.InputRequiredResult | InputRequiredToolResult): + raise MCPError( + code=mcp_types.INTERNAL_ERROR, + message=( + f"Tool {tool.name!r} requested input while running as a background " + "task. Input-required (multi-round-trip) tools need a live request " + "to answer the prompt and cannot run as tasks." + ), + ) + mcp_result = tool.convert_result(raw_value).to_mcp_result() + if isinstance(mcp_result, mcp_types.CallToolResult): + call_tool_result = mcp_result + elif isinstance(mcp_result, tuple): + content, structured_content = mcp_result + call_tool_result = mcp_types.CallToolResult( + content=content, structured_content=structured_content + ) + else: + call_tool_result = mcp_types.CallToolResult(content=mcp_result) + return call_tool_result.model_dump(by_alias=True, mode="json", exclude_none=True) + + +async def tasks_get(server: FastMCP, task_id: str) -> GetTaskResult: + """Handle ``tasks/get``: the detailed task with its result/error/inputs inlined.""" + docket = server._docket + if docket is None: + raise _task_not_found(task_id) + + task_scope = get_task_scope() + execution, task_key, created_at, poll_interval_ms = await _lookup_task( + docket, task_scope, task_id + ) + await execution.sync() + + created_at_iso = _normalize_iso_timestamp(created_at) + now_iso = datetime.now(timezone.utc).isoformat() + ttl_ms = _ttl_ms(docket) + + def build( + status: Literal[ + "working", "input_required", "completed", "failed", "cancelled" + ], + **payload: Any, + ) -> GetTaskResult: + return GetTaskResult( + task_id=task_id, + status=status, + created_at=created_at_iso, + last_updated_at=now_iso, + ttl_ms=ttl_ms, + poll_interval_ms=poll_interval_ms, + **payload, + ) + + # An outstanding input request outranks the Docket "running" state: the task + # is parked in the worker waiting for tasks/update, so it is input_required. + if execution.state in _WORKING_STATES: + outstanding = await read_outstanding_inputs(docket, task_scope, task_id) + if outstanding: + return build("input_required", input_requests=outstanding) + + if execution.state == ExecutionState.COMPLETED: + raw_value = await execution.get_result(timeout=timedelta(seconds=0)) + tool = await _resolve_tool(server, task_key) + return build("completed", result=_inline_result(tool, raw_value)) + + if execution.state == ExecutionState.FAILED: + message = "Task failed" + try: + await execution.get_result(timeout=timedelta(seconds=0)) + # On a FAILED execution, get_result re-raises the exception the task + # itself raised — an arbitrary user-defined type, so no narrower catch + # exists. Its message becomes the task's error payload. + except Exception as error: + message = str(error) + return build( + "failed", + status_message=message, + error={"code": mcp_types.INTERNAL_ERROR, "message": message}, + ) + + if execution.state == ExecutionState.CANCELLED: + return build("cancelled") + + status_message = None + if execution.progress and execution.progress.message: + status_message = execution.progress.message + return build("working", status_message=status_message) + + +async def tasks_update( + server: FastMCP, task_id: str, input_responses: dict[str, Any] +) -> UpdateTaskResult: + """Handle ``tasks/update``: deliver input responses to the parked worker.""" + docket = server._docket + if docket is None: + raise _task_not_found(task_id) + + task_scope = get_task_scope() + # Resolve within scope so a cross-scope update is a "not found", not a no-op. + await _lookup_task(docket, task_scope, task_id) + await deliver_input_responses(docket, task_scope, task_id, input_responses) + return UpdateTaskResult() + + +async def tasks_cancel(server: FastMCP, task_id: str) -> CancelTaskResult: + """Handle ``tasks/cancel``: cooperatively cancel the task, empty ack.""" + docket = server._docket + if docket is None: + raise _task_not_found(task_id) + + task_scope = get_task_scope() + execution, _task_key, _created_at, _poll = await _lookup_task( + docket, task_scope, task_id + ) + await docket.cancel(execution.key) + return CancelTaskResult() diff --git a/fastmcp_tasks/fastmcp_tasks/input_store.py b/fastmcp_tasks/fastmcp_tasks/input_store.py new file mode 100644 index 000000000..5c032be06 --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/input_store.py @@ -0,0 +1,169 @@ +"""In-task input store for SEP-2663 poll-based elicitation. + +When a background task calls ``ctx.elicit()`` it has no live request to carry the +prompt. SEP-2663 handles this by polling: the worker parks an *input request* +here, the task's ``tasks/get`` status flips to ``input_required`` with the +outstanding requests, the client answers via ``tasks/update``, and the parked +worker resumes. + +This is the reworked SEP-1686 elicitation module. The Redis request/response +mechanics — a per-request hash the poll surface reads and a per-key list the +worker blocks on with ``BLPOP`` — are preserved. What's gone is the *push +envelope*: the old code sent a ``notifications/tasks/status`` through the +distributed notification queue to wake the client. Under SEP-2663 the client +discovers the outstanding request by polling ``tasks/get``, so no push is needed. +""" + +from __future__ import annotations + +import json +import logging +from typing import TYPE_CHECKING, Any + +import mcp_types +from redis.exceptions import RedisError + +from fastmcp_tasks.context import get_task_context +from fastmcp_tasks.keys import task_redis_prefix + +if TYPE_CHECKING: + from docket import Docket + + from fastmcp.server.context import Context + +logger = logging.getLogger(__name__) + +# How long a parked input request (and any delivered response) lives before +# expiring. A task blocked on input holds a worker slot, so this doubles as the +# maximum time a worker waits for the client to answer. +INPUT_TTL_SECONDS = 3600 + + +def _requests_key(docket: Docket, task_scope: str | None, task_id: str) -> str: + """Redis hash of outstanding input requests, keyed by input key.""" + return docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:input:requests") + + +def _response_key( + docket: Docket, task_scope: str | None, task_id: str, input_key: str +) -> str: + """Redis list the worker blocks on for a single input key's response.""" + return docket.key( + f"{task_redis_prefix(task_scope)}:{task_id}:input:resp:{input_key}" + ) + + +def _elicitation_input_request(message: str, schema: dict[str, Any]) -> dict[str, Any]: + """Build the SEP-2663 ``InputRequest`` for an elicitation (an ElicitRequest).""" + return { + "method": "elicitation/create", + "params": {"message": message, "requestedSchema": schema}, + } + + +async def elicit_in_task( + context: Context, message: str, schema: dict[str, Any] +) -> mcp_types.ElicitResult: + """Park an elicitation request and block until the client answers it. + + Installed as core's in-task elicitation handler by ``TasksExtension``. Parks + an input request keyed by the task's own id (one outstanding elicitation per + task at a time — the polling model is inherently sequential), flips the + task's polled status to ``input_required``, and blocks on the response list. + Returns the client's ``ElicitResult``; on timeout or a missing task context, + returns a ``cancel`` action so the worker never hangs indefinitely. + """ + task_context = get_task_context() + if task_context is None: + logger.warning("elicit_in_task called outside a task worker; cancelling") + return mcp_types.ElicitResult(action="cancel", content=None) + + docket = context.fastmcp._docket + if docket is None: + from fastmcp_tasks.dependencies import _current_docket + + docket = _current_docket.get() + if docket is None: + return mcp_types.ElicitResult(action="cancel", content=None) + + task_scope = task_context.task_scope + task_id = task_context.task_id + # One elicitation outstanding per task: key the request by the task id so the + # inputRequests map surfaced by tasks/get is stable and answerable. + input_key = task_id + + requests_key = _requests_key(docket, task_scope, task_id) + response_key = _response_key(docket, task_scope, task_id, input_key) + request_payload = _elicitation_input_request(message, schema) + + async with docket.redis() as redis: + await redis.hset(requests_key, input_key, json.dumps(request_payload)) + await redis.expire(requests_key, INPUT_TTL_SECONDS) + + try: + async with docket.redis() as redis: + result = await redis.blpop([response_key], timeout=INPUT_TTL_SECONDS) + except (RedisError, OSError) as exc: + logger.warning("BLPOP failed for task %s input; cancelling: %s", task_id, exc) + result = None + + async with docket.redis() as redis: + await redis.hdel(requests_key, input_key) + await redis.delete(response_key) + + if not result: + return mcp_types.ElicitResult(action="cancel", content=None) + + _key, raw = result + response = json.loads(raw) + return mcp_types.ElicitResult( + action=response.get("action", "accept"), + content=response.get("content"), + ) + + +async def read_outstanding_inputs( + docket: Docket, task_scope: str | None, task_id: str +) -> dict[str, Any]: + """Return the task's outstanding input requests, keyed by input key. + + Empty when the task is not waiting on input. Consumed by ``tasks/get`` to + build the ``input_required`` status and its ``inputRequests`` snapshot. + """ + requests_key = _requests_key(docket, task_scope, task_id) + async with docket.redis() as redis: + raw = await redis.hgetall(requests_key) + outstanding: dict[str, Any] = {} + for key, value in raw.items(): + key_str = key.decode() if isinstance(key, bytes) else key + value_str = value.decode() if isinstance(value, bytes) else value + try: + outstanding[key_str] = json.loads(value_str) + except json.JSONDecodeError: + continue + return outstanding + + +async def deliver_input_responses( + docket: Docket, + task_scope: str | None, + task_id: str, + responses: dict[str, Any], +) -> None: + """Deliver ``tasks/update`` responses to the parked worker(s). + + For each response whose key names an outstanding request, pushes the + response onto that key's list (waking the worker's ``BLPOP``) and removes the + request. Responses for unknown or already-satisfied keys are ignored, as the + spec requires. + """ + requests_key = _requests_key(docket, task_scope, task_id) + async with docket.redis() as redis: + for input_key, response in responses.items(): + outstanding = await redis.hget(requests_key, input_key) + if outstanding is None: + continue + response_key = _response_key(docket, task_scope, task_id, input_key) + await redis.rpush(response_key, json.dumps(response)) + await redis.expire(response_key, INPUT_TTL_SECONDS) + await redis.hdel(requests_key, input_key) diff --git a/fastmcp_tasks/fastmcp_tasks/lifespan.py b/fastmcp_tasks/fastmcp_tasks/lifespan.py index 45df1a7c5..925bd6325 100644 --- a/fastmcp_tasks/fastmcp_tasks/lifespan.py +++ b/fastmcp_tasks/fastmcp_tasks/lifespan.py @@ -1,19 +1,17 @@ """Docket lifecycle for FastMCP background tasks. -Extracted from ``fastmcp.server.mixins.lifespan.LifespanMixin._docket_lifespan`` -during the SEP-1686 -> SEP-2663 migration. The logic — start Docket and a Worker -at the runtime-tree root when there are task-enabled components, register those -components' callables, and run the worker with the snapshot-restore dependency — -is preserved verbatim so Phase 3 can drive it from ``TasksExtension.lifespan()``. - -Nothing in core calls this after Phase 2; it is engine code parked here for the -Phase 3 adapter. +Extracted from the SEP-1686 ``LifespanMixin._docket_lifespan`` and driven by +``TasksExtension.lifespan()``. Core's ``_extensions_lifespan`` already enters +this once per runtime tree at the root and defers on mounted children, and +``SharedContext`` plus the server ContextVar are established before extension +lifespans run — so this no longer manages either. It starts Docket and a Worker +when there are task-enabled components, registers those components' callables, +and runs the worker (with the snapshot-restore dependency) until shutdown. """ from __future__ import annotations import asyncio -import weakref from collections.abc import AsyncIterator from contextlib import asynccontextmanager, suppress from typing import TYPE_CHECKING, Any @@ -22,26 +20,25 @@ from fastmcp.utilities.logging import get_logger if TYPE_CHECKING: from fastmcp.server.server import FastMCP + from fastmcp_tasks.settings import DocketSettings logger = get_logger(__name__) @asynccontextmanager -async def docket_lifespan(server: FastMCP) -> AsyncIterator[None]: +async def docket_lifespan( + server: FastMCP, settings: DocketSettings +) -> AsyncIterator[None]: """Manage the Docket instance and Worker for background task execution. - Docket infrastructure is only initialized if: - 1. pydocket is installed (fastmcp[tasks] extra) - 2. There are task-enabled components (task_config.mode != 'forbidden') - Sets ``server._docket`` / ``server._worker`` for the duration and registers - each task-enabled component's callable with the Docket, then runs the worker - until the context exits. + each task-enabled component's callable, then runs the worker until the + context exits. A no-op if pydocket is unavailable or the server declares no + task-enabled components. """ from docket import Depends, Docket, Worker import fastmcp - from fastmcp.server.dependencies import _current_server from fastmcp_tasks.components import register_component_with_docket from fastmcp_tasks.context import restore_task_snapshot from fastmcp_tasks.dependencies import ( @@ -49,78 +46,60 @@ async def docket_lifespan(server: FastMCP) -> AsyncIterator[None]: _current_worker, is_docket_available, ) - from fastmcp_tasks.settings import DocketSettings - docket_settings = DocketSettings() - - # Set FastMCP server in ContextVar so CurrentFastMCP can access it - # (use weakref to avoid reference cycles) - server_token = _current_server.set(weakref.ref(server)) + if not is_docket_available(): + yield + return try: - if not is_docket_available(): - yield - return + candidates = list(await server.get_tasks()) + except Exception as e: + logger.warning(f"Failed to collect task components: {e}") + if fastmcp.settings.mounted_components_raise_on_load_error: + raise + candidates = [] - # Collect task-enabled components at startup with all transforms applied. - # Components must be available now to be registered with Docket workers; - # dynamically added components after startup won't be registered. + # get_tasks() applies server-level transforms that can inject non-task tools; + # re-filter by the actual task config (the recorded landmine). + task_components = [c for c in candidates if c.task_config.supports_tasks()] + if not task_components: + yield + return + + async with Docket(name=settings.name, url=settings.url) as docket: + server._docket = docket + for component in task_components: + register_component_with_docket(component, docket) + + docket_token = _current_docket.set(docket) try: - task_components = list(await server.get_tasks()) - except Exception as e: - logger.warning(f"Failed to get tasks: {e}") - if fastmcp.settings.mounted_components_raise_on_load_error: - raise - task_components = [] + worker_kwargs: dict[str, Any] = { + "concurrency": settings.concurrency, + "redelivery_timeout": settings.redelivery_timeout, + "reconnection_delay": settings.reconnection_delay, + "minimum_check_interval": settings.minimum_check_interval, + } + if settings.worker_name: + worker_kwargs["name"] = settings.worker_name - if not task_components: - yield - return - - async with Docket( - name=docket_settings.name, - url=docket_settings.url, - ) as docket: - server._docket = docket - - for component in task_components: - register_component_with_docket(component, docket) - - docket_token = _current_docket.set(docket) - try: - worker_kwargs: dict[str, Any] = { - "concurrency": docket_settings.concurrency, - "redelivery_timeout": docket_settings.redelivery_timeout, - "reconnection_delay": docket_settings.reconnection_delay, - "minimum_check_interval": docket_settings.minimum_check_interval, - } - if docket_settings.worker_name: - worker_kwargs["name"] = docket_settings.worker_name - - # Create and start Worker. The restore_task_snapshot worker-level - # dependency runs before every task so the per-task snapshot - # ContextVar is populated before user code or task-scoped - # dependencies observe it. - async with Worker( - docket, - dependencies=[Depends(restore_task_snapshot)], - **worker_kwargs, - ) as worker: - server._worker = worker - worker_token = _current_worker.set(worker) + async with Worker( + docket, + dependencies=[Depends(restore_task_snapshot)], + **worker_kwargs, + ) as worker: + server._worker = worker + worker_token = _current_worker.set(worker) + try: + worker_task = asyncio.create_task(worker.run_forever()) try: - worker_task = asyncio.create_task(worker.run_forever()) - try: - yield - finally: - worker_task.cancel() - with suppress(asyncio.CancelledError): - await worker_task + yield finally: - _current_worker.reset(worker_token) - server._worker = None - finally: - _current_docket.reset(docket_token) - server._docket = None - finally: - _current_server.reset(server_token) + worker_task.cancel() + with suppress(asyncio.CancelledError): + await worker_task + finally: + _current_worker.reset(worker_token) + server._worker = None + finally: + _current_docket.reset(docket_token) + server._docket = None diff --git a/fastmcp_tasks/fastmcp_tasks/models.py b/fastmcp_tasks/fastmcp_tasks/models.py new file mode 100644 index 000000000..d6b3a9875 --- /dev/null +++ b/fastmcp_tasks/fastmcp_tasks/models.py @@ -0,0 +1,167 @@ +"""SEP-2663 tasks-extension wire models. + +The `io.modelcontextprotocol/tasks` extension (SEP-2663) defines its own wire +shapes, distinct from the SEP-1686 task types the MCP SDK still ships +(`mcp_types.Task` uses `ttl`/`pollInterval`; SEP-2663 uses `ttlMs`/`pollIntervalMs` +and a *flat* `CreateTaskResult` rather than a nested `{task: ...}`). These models +serialize to the SEP-2663 shapes and are validated against the vendored draft +JSON schema in the test suite. + +A note on `_meta`: the draft schema composes result shapes as +`allOf[Result, Task]`, and the `Task` arm carries `additionalProperties: false` +without listing `_meta`. A `_meta` key therefore fails schema validation on those +results. These models leave `_meta` unset and rely on the runner's +`exclude_none=True` dump to omit it, so serialized instances validate cleanly. +`ttlMs` is required-but-nullable in the schema; in practice the engine always +emits a numeric value (Docket carries a default execution TTL), so the +`exclude_none` dump never drops it. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from mcp_types import RequestParams, Result +from pydantic import BaseModel, ConfigDict, Field + +__all__ = [ + "MISSING_REQUIRED_CLIENT_CAPABILITY", + "TaskStatus", + "CreateTaskResult", + "GetTaskResult", + "UpdateTaskResult", + "CancelTaskResult", + "GetTaskParams", + "UpdateTaskParams", + "CancelTaskParams", + "GetTaskRequest", + "UpdateTaskRequest", + "CancelTaskRequest", + "missing_capability_error_data", +] + +#: JSON-RPC error code for "Missing Required Client Capability" (SEP-2663). A +#: tool whose task mode is `required` returns this when the client did not opt +#: the tasks extension in for the request. +MISSING_REQUIRED_CLIENT_CAPABILITY = -32003 + +TaskStatus = Literal[ + "working", "input_required", "completed", "failed", "cancelled" +] + + +class _TaskFields(BaseModel): + """The flat task fields shared by every SEP-2663 task result shape. + + Serializes to the schema's `Task` object (camelCase aliases, `ttlMs` + required-but-nullable). No `_meta`: the schema's `additionalProperties: + false` on the task arm forbids it (see module docstring). + """ + + # Serialization aliases only: these result models are constructed by field + # name (the engine builds them) and dumped to camelCase by the runner + # (`model_dump(by_alias=True)`). Wire *validation* of results is the client's + # concern. + model_config = ConfigDict(populate_by_name=True) + + task_id: str = Field(serialization_alias="taskId") + status: TaskStatus + created_at: str = Field(serialization_alias="createdAt") + last_updated_at: str = Field(serialization_alias="lastUpdatedAt") + ttl_ms: float | None = Field(serialization_alias="ttlMs") + status_message: str | None = Field(default=None, serialization_alias="statusMessage") + poll_interval_ms: float | None = Field( + default=None, serialization_alias="pollIntervalMs" + ) + + +class CreateTaskResult(_TaskFields): + """Result of an augmented `tools/call` that the server ran as a task. + + A flat merge of `Result` and `Task` (SEP-2663): the finished task stub the + client polls with `tasks/get`. Status is typically `working`. + """ + + +class GetTaskResult(_TaskFields): + """Result of `tasks/get`: the detailed task (`Result & DetailedTask`). + + Carries exactly one of `result` (completed), `error` (failed), or + `input_requests` (input_required) alongside the flat task fields, matching + the schema's 5-status union. The three payload fields default to `None` and + are dropped from the wire dump for the statuses that do not use them. + """ + + result: dict[str, Any] | None = None + error: dict[str, Any] | None = None + input_requests: dict[str, Any] | None = Field( + default=None, serialization_alias="inputRequests" + ) + + +class UpdateTaskResult(Result): + """Empty acknowledgement for `tasks/update` (SEP-2663 `Result`).""" + + +class CancelTaskResult(Result): + """Empty acknowledgement for `tasks/cancel` (SEP-2663 `Result`).""" + + +class GetTaskParams(RequestParams): + """Params for `tasks/get` / `tasks/cancel`: the target task id.""" + + model_config = ConfigDict(populate_by_name=True) + + task_id: str = Field(alias="taskId") + + +# `tasks/cancel` params are identical to `tasks/get` (just `taskId`). +CancelTaskParams = GetTaskParams + + +class UpdateTaskParams(RequestParams): + """Params for `tasks/update`: task id plus the caller's input responses.""" + + model_config = ConfigDict(populate_by_name=True) + + task_id: str = Field(alias="taskId") + input_responses: dict[str, Any] = Field(alias="inputResponses") + + +class GetTaskRequest(BaseModel): + """`tasks/get` request envelope (used by tests and clients).""" + + model_config = ConfigDict(populate_by_name=True) + + method: Literal["tasks/get"] = "tasks/get" + params: GetTaskParams + + +class UpdateTaskRequest(BaseModel): + """`tasks/update` request envelope.""" + + model_config = ConfigDict(populate_by_name=True) + + method: Literal["tasks/update"] = "tasks/update" + params: UpdateTaskParams + + +class CancelTaskRequest(BaseModel): + """`tasks/cancel` request envelope.""" + + model_config = ConfigDict(populate_by_name=True) + + method: Literal["tasks/cancel"] = "tasks/cancel" + params: GetTaskParams + + +def missing_capability_error_data() -> dict[str, Any]: + """Build the `data.requiredCapabilities` payload for a -32003 error. + + A `required`-mode tool called without the client opting the tasks extension + in for the request returns this so the client learns which capability to + declare. + """ + from fastmcp.utilities.tasks import TASKS_EXTENSION_ID + + return {"requiredCapabilities": {"extensions": {TASKS_EXTENSION_ID: {}}}} diff --git a/fastmcp_tasks/fastmcp_tasks/settings.py b/fastmcp_tasks/fastmcp_tasks/settings.py index 57b6970a4..4f9fec622 100644 --- a/fastmcp_tasks/fastmcp_tasks/settings.py +++ b/fastmcp_tasks/fastmcp_tasks/settings.py @@ -2,8 +2,8 @@ Moved out of ``fastmcp.settings`` during the SEP-1686 -> SEP-2663 migration. The ``FASTMCP_DOCKET_*`` environment prefix is unchanged so existing -deployments keep working. Phase 3 wires this configuration into -``TasksExtension``. +deployments keep working. ``TasksExtension`` reads this configuration (its +constructor overrides the env defaults). """ from __future__ import annotations diff --git a/fastmcp_tasks/fastmcp_tasks/worker_cli.py b/fastmcp_tasks/fastmcp_tasks/worker_cli.py index d39ddfce8..01af894b4 100644 --- a/fastmcp_tasks/fastmcp_tasks/worker_cli.py +++ b/fastmcp_tasks/fastmcp_tasks/worker_cli.py @@ -105,3 +105,9 @@ def worker( except KeyboardInterrupt: console.print("\n[yellow]Worker stopped[/yellow]") sys.exit(0) + + +if __name__ == "__main__": + # Enables `python -m fastmcp_tasks.worker_cli worker ` for running an + # out-of-process worker now that core dropped the `fastmcp tasks` subcommand. + tasks_app() diff --git a/pyproject.toml b/pyproject.toml index 46174c506..e4deaae1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -155,10 +155,9 @@ exclude = [ "examples/providers/sqlite", # needs aiosqlite "examples/memory.py", # needs asyncpg, numpy, pydantic_ai, pgvector "examples/get_file.py", # needs aiohttp - # Dormant SEP-1686 task tests: skipped at runtime pending the Phase 3 - # TasksExtension (SEP-2663). They reference task APIs that are removed from - # core and return in the fastmcp-tasks extension, so they don't type-check - # against core until then. Drop this exclusion when Phase 3 lands. + # The moved task tests pass at runtime but carry ty diagnostics (mostly + # None-narrowing on optional result fields); a follow-up commit fixes them + # and removes this exclusion. "tests/tasks", ] diff --git a/tests/cli/test_tasks.py b/tests/cli/test_tasks.py index a2ea42ede..fc0fcdaa1 100644 --- a/tests/cli/test_tasks.py +++ b/tests/cli/test_tasks.py @@ -1,31 +1,30 @@ """Tests for the fastmcp tasks CLI.""" import pytest +from fastmcp_tasks.settings import docket_settings from fastmcp_tasks.worker_cli import check_distributed_backend, tasks_app -from fastmcp.utilities.tests import temporary_settings - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) - class TestCheckDistributedBackend: """Test the distributed backend checker function.""" - def test_succeeds_with_redis_url(self): + def test_succeeds_with_redis_url(self, monkeypatch: pytest.MonkeyPatch): """Test that it succeeds with Redis URL.""" - with temporary_settings(docket__url="redis://localhost:6379/0"): + # Docket settings moved to `fastmcp_tasks.settings.DocketSettings` + # (env prefix `FASTMCP_DOCKET_`), so patch the settings object directly. + monkeypatch.setattr(docket_settings, "url", "redis://localhost:6379/0") + check_distributed_backend() + + def test_exits_with_helpful_error_for_memory_url( + self, monkeypatch: pytest.MonkeyPatch + ): + """Test that it exits with helpful error for memory:// URLs.""" + monkeypatch.setattr(docket_settings, "url", "memory://test-123") + with pytest.raises(SystemExit) as exc_info: check_distributed_backend() - def test_exits_with_helpful_error_for_memory_url(self): - """Test that it exits with helpful error for memory:// URLs.""" - with temporary_settings(docket__url="memory://test-123"): - with pytest.raises(SystemExit) as exc_info: - check_distributed_backend() - - assert isinstance(exc_info.value, SystemExit) - assert exc_info.value.code == 1 + assert isinstance(exc_info.value, SystemExit) + assert exc_info.value.code == 1 class TestWorkerCommand: diff --git a/tests/client/client/test_client.py b/tests/client/client/test_client.py index bfe74f68c..38f733f63 100644 --- a/tests/client/client/test_client.py +++ b/tests/client/client/test_client.py @@ -886,7 +886,7 @@ async def test_client_list_dict_return_type(): assert result.data == [{"city": "NYC", "temp": 72}, {"city": "LA", "temp": 85}] -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def test_client_new_resets_mutable_task_state(fastmcp_server): """Client.new() should not share mutable task tracking structures.""" client = Client(transport=FastMCPTransport(fastmcp_server)) @@ -903,7 +903,7 @@ def test_client_new_resets_mutable_task_state(fastmcp_server): assert clone._submitted_task_ids is not client._submitted_task_ids # ty: ignore -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def test_client_new_rebinds_default_task_notification_handler(fastmcp_server): """Client.new() should bind the default task handler to the cloned client.""" client = Client(transport=FastMCPTransport(fastmcp_server)) diff --git a/tests/client/telemetry/test_client_task_tracing.py b/tests/client/telemetry/test_client_task_tracing.py index ad5081e7d..4c93a2fed 100644 --- a/tests/client/telemetry/test_client_task_tracing.py +++ b/tests/client/telemetry/test_client_task_tracing.py @@ -10,9 +10,7 @@ from opentelemetry.trace import SpanKind from fastmcp import Client, FastMCP -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def assert_propagating_client_span( diff --git a/tests/client/test_client_extensions.py b/tests/client/test_client_extensions.py index 8681fc817..03af1e89c 100644 --- a/tests/client/test_client_extensions.py +++ b/tests/client/test_client_extensions.py @@ -142,7 +142,7 @@ def test_extension_populates_claim_by_model_index(): assert client._claim_by_model[ClaimedResult].result_type == CLAIMED_TYPE -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def test_binding_composes_with_internal_task_binding(): """User binding is appended to (not replacing) the task-status binding.""" client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) @@ -154,7 +154,7 @@ def test_binding_composes_with_internal_task_binding(): assert methods[0] == TASK_STATUS_METHOD -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def test_no_extensions_leaves_only_task_binding(): """Without extensions, only the internal task-status binding is registered.""" client = Client(FastMCP("srv")) @@ -165,7 +165,7 @@ def test_no_extensions_leaves_only_task_binding(): assert client._claim_by_model == {} -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") def test_new_preserves_extension_composition(): """new() rebuilds the clone with both the task binding and user bindings.""" client = Client(FastMCP("srv"), extensions=[_DemoExtension()]) @@ -207,7 +207,7 @@ def test_result_claims_merge_with_extension_claims(): assert set(client._claim_by_model) == {ClaimedResult, ExtraClaimed} -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") async def test_user_binding_clobbering_task_method_is_rejected(): """A user extension binding the task-status method cannot silently replace it. @@ -237,7 +237,7 @@ async def test_user_binding_clobbering_task_method_is_rejected(): pass -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") async def test_both_bindings_fire_against_live_server(): """The internal task binding and a user extension binding both fire. diff --git a/tests/client/transports/test_memory_transport.py b/tests/client/transports/test_memory_transport.py index 5bbe3563e..5153ecbf4 100644 --- a/tests/client/transports/test_memory_transport.py +++ b/tests/client/transports/test_memory_transport.py @@ -7,9 +7,12 @@ Client(server) with an in-process FastMCP server. import time import pytest +from docket import Docket from fastmcp import Client, FastMCP from fastmcp.client.transports import FastMCPTransport +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import submit_task, wait_for_task def test_transport_repr_includes_server_name(): @@ -18,9 +21,18 @@ def test_transport_repr_includes_server_name(): assert repr(transport) == "" -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") +@pytest.fixture +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to this test's loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + @pytest.mark.timeout(10) -async def test_task_teardown_does_not_hang(): +async def test_task_teardown_does_not_hang(reset_docket_memory_server): """In-memory transport must tear down in under 2 seconds after a task call. This is a regression test for a teardown ordering bug where the Docket @@ -39,8 +51,13 @@ async def test_task_teardown_does_not_hang(): If this test takes ~5 seconds, the context manager nesting in FastMCPTransport.connect_session() has been reversed — the lifespan must be the OUTER context and the task group must be the INNER context. + + There is no client task-submission API yet (Phase 4), so the task is + driven server-side within the live in-memory session; the teardown path + being exercised is the same either way. """ mcp = FastMCP("teardown-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def fast_tool(x: int) -> int: @@ -48,10 +65,12 @@ async def test_task_teardown_does_not_hang(): t0 = time.monotonic() - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("fast_tool", {"x": 21}, task=True) - result = await task.result() - assert result.data == 42 + async with Client(mcp): + created = await submit_task(mcp, "fast_tool", {"x": 21}) + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result is not None + assert final.result["structuredContent"] == {"result": 42} elapsed = time.monotonic() - t0 diff --git a/tests/server/http/test_http_dependencies.py b/tests/server/http/test_http_dependencies.py index 531b8d1f6..704baebc3 100644 --- a/tests/server/http/test_http_dependencies.py +++ b/tests/server/http/test_http_dependencies.py @@ -1,12 +1,43 @@ import json import pytest +from docket import Docket +from fastmcp_tasks.context import _recall_snapshot, get_task_context from mcp_types import TextContent, TextResourceContents from starlette.requests import Request -from fastmcp.server.dependencies import CurrentHeaders, CurrentRequest, get_http_request +from fastmcp.server.dependencies import get_http_request +from fastmcp.server.http import _current_http_request from fastmcp.server.server import FastMCP from fastmcp.utilities.tests import ASGIServer, asgi_server +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import running_task_server, submit_task, wait_for_task + + +@pytest.fixture +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to this test's event loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + +def _http_request_with_headers(headers: dict[str, str]) -> Request: + """Build a minimal Starlette HTTP request carrying the given headers.""" + raw_headers = [(k.lower().encode(), v.encode()) for k, v in headers.items()] + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": raw_headers, + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "client": ("testclient", 12345), + } + return Request(scope) def fastmcp_server(): @@ -146,50 +177,80 @@ async def test_get_http_headers_excludes_content_type(sse_server: ASGIServer): assert headers["x-custom-header"] == "should-be-included" -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") -async def test_background_task_can_read_snapshotted_request_headers(): - """Background tools can still access request headers via get_http_request().""" +def _worker_snapshot_headers() -> dict[str, str]: + """Read the HTTP headers snapshotted at task submission from inside a worker.""" + task_info = get_task_context() + snapshot = _recall_snapshot(task_info.task_id) if task_info is not None else None + if snapshot is None or snapshot.http_headers is None: + return {} + return dict(snapshot.http_headers) + + +async def test_background_task_can_read_snapshotted_request_headers( + reset_docket_memory_server, +): + """A background task worker reads the HTTP headers snapshotted at submission. + + There is no client task-submission API yet (Phase 4), so the task is driven + in-process: an HTTP request is bound while the task is submitted, and the + worker reads the request headers back from the restored task-context + snapshot. + """ server = FastMCP() + server.add_extension(TasksExtension()) @server.tool(task=True) async def check_request_header() -> str: - request = get_http_request() - return request.headers.get("x-tenant-id", "missing") + return _worker_snapshot_headers().get("x-tenant-id", "missing") - async with asgi_server(server, transport="sse") as running_server: - async with running_server.client( - headers={"X-Tenant-ID": "tenant-123"} - ) as client: - task = await client.call_tool("check_request_header", task=True) - result = await task.result() - assert result.data == "tenant-123" + request = _http_request_with_headers({"X-Tenant-ID": "tenant-123"}) + async with running_task_server(server): + token = _current_http_request.set(request) + try: + created = await submit_task(server, "check_request_header", {}) + finally: + _current_http_request.reset(token) + + final = await wait_for_task(server, created.task_id) + + assert final.status == "completed" + assert final.result is not None + assert final.result["structuredContent"] == {"result": "tenant-123"} -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") -async def test_background_task_current_http_dependencies_restore_headers(): - """CurrentHeaders/CurrentRequest work in task workers without explicit Context.""" +async def test_background_task_snapshot_preserves_all_request_headers( + reset_docket_memory_server, +): + """The task snapshot preserves every request header, including authorization.""" server = FastMCP() + server.add_extension(TasksExtension()) @server.tool(task=True) - async def check_headers( - headers: dict[str, str] = CurrentHeaders(), - request: Request = CurrentRequest(), - ) -> dict[str, str]: + async def check_headers() -> dict[str, str]: + headers = _worker_snapshot_headers() return { "authorization": headers.get("authorization", "missing"), - "tenant": request.headers.get("x-tenant-id", "missing"), + "tenant": headers.get("x-tenant-id", "missing"), } - async with asgi_server(server, transport="sse") as running_server: - async with running_server.client( - headers={ - "Authorization": "Bearer tenant-token", - "X-Tenant-ID": "tenant-456", - } - ) as client: - task = await client.call_tool("check_headers", task=True) - result = await task.result() - assert result.data == { - "authorization": "Bearer tenant-token", - "tenant": "tenant-456", - } + request = _http_request_with_headers( + { + "Authorization": "Bearer tenant-token", + "X-Tenant-ID": "tenant-456", + } + ) + async with running_task_server(server): + token = _current_http_request.set(request) + try: + created = await submit_task(server, "check_headers", {}) + finally: + _current_http_request.reset(token) + + final = await wait_for_task(server, created.task_id) + + assert final.status == "completed" + assert final.result is not None + assert final.result["structuredContent"] == { + "authorization": "Bearer tenant-token", + "tenant": "tenant-456", + } diff --git a/tests/server/mount/test_advanced.py b/tests/server/mount/test_advanced.py index 89f590e63..24c9765cd 100644 --- a/tests/server/mount/test_advanced.py +++ b/tests/server/mount/test_advanced.py @@ -1,6 +1,7 @@ """Advanced mounting scenarios.""" import pytest +from docket import Docket from mcp_types import TextContent from starlette.routing import Route @@ -8,6 +9,18 @@ from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.server.providers import FastMCPProvider from fastmcp.server.providers.wrapped_provider import _WrappedProvider +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import running_task_server + + +@pytest.fixture +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to this test's event loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") class TestDynamicChanges: @@ -598,17 +611,19 @@ class TestMountedServerDocketBehavior: includes Docket creation. """ - @pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") - async def test_mounted_server_does_not_have_docket(self): + async def test_mounted_server_does_not_have_docket( + self, reset_docket_memory_server + ): """Test that a mounted server doesn't create its own Docket. MountedProvider.lifespan() should call only the server's _lifespan (user-defined lifespan), not _lifespan_manager (which includes Docket). """ main_app = FastMCP("MainApp") + main_app.add_extension(TasksExtension()) sub_app = FastMCP("SubApp") - # Need a task-enabled component to trigger Docket initialization + # A task-enabled component on the parent makes it own a Docket. @main_app.tool(task=True) async def _trigger_docket() -> str: return "trigger" @@ -619,21 +634,15 @@ class TestMountedServerDocketBehavior: main_app.mount(sub_app, "sub") - # After running the main app's lifespan, the sub app should not have - # its own Docket instance - async with Client(main_app) as client: - # The main app should have a docket (created by _lifespan_manager) - # because it has a task-enabled component + # After entering the parent's lifespan, only the parent owns a Docket. + async with running_task_server(main_app): + # The parent owns a Docket because it has a task-enabled component. assert main_app.docket is not None - # The mounted sub app should NOT have its own docket - # It uses the parent's docket for background tasks + # The mounted child does NOT own its own Docket; it uses the + # parent's Docket for background tasks. assert sub_app.docket is None - # But the tool should still work (prefixed as sub_my_tool) - result = await client.call_tool("sub_my_tool", {}) - assert result.data == "test" - class TestComponentServicePrefixLess: """Test that enable/disable works with prefix-less mounted servers.""" diff --git a/tests/server/test_dependencies.py b/tests/server/test_dependencies.py index 1916c3428..cef7cbf45 100644 --- a/tests/server/test_dependencies.py +++ b/tests/server/test_dependencies.py @@ -3,17 +3,29 @@ from contextlib import asynccontextmanager, contextmanager import pytest +from docket import Docket from mcp_types import TextContent, TextResourceContents from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.dependencies import CurrentContext, Depends, Shared from fastmcp.server.context import Context +from fastmcp_tasks import TasksExtension from tests.conftest import make_server_request_context HUZZAH = "huzzah!" +@pytest.fixture +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to this test's event loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + class Connection: """Test connection that tracks whether it's currently open.""" @@ -1193,8 +1205,9 @@ class TestSharedDependencies: ) assert call_count == 1 - @pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") - async def test_shared_resolves_on_task_capable_server(self): + async def test_shared_resolves_on_task_capable_server( + self, reset_docket_memory_server + ): """Shared() dependencies resolve on a normal request even when the server has task-enabled components. @@ -1205,6 +1218,7 @@ class TestSharedDependencies: on ordinary (non-task) calls. """ mcp = FastMCP("task-capable-server") + mcp.add_extension(TasksExtension()) call_count = 0 diff --git a/tests/server/test_extensions.py b/tests/server/test_extensions.py index a3e65ad62..c112f67b1 100644 --- a/tests/server/test_extensions.py +++ b/tests/server/test_extensions.py @@ -472,6 +472,19 @@ async def test_duplicate_identifier_rejected(): mcp.add_extension(Ext()) +async def test_registration_after_lifespan_start_rejected(): + """Registering once the server is serving would skip the extension's + lifespan, leaving it silently half-active — so it raises instead.""" + + class Ext(ServerExtension): + identifier = EXT_ID + + mcp = FastMCP("t") + async with Client(mcp, mode="auto"): + with pytest.raises(RuntimeError, match="lifespan has already started"): + mcp.add_extension(Ext()) + + def test_spec_method_name_rejected(): async def handler(ctx: Any, params: Any) -> None: return None diff --git a/tests/server/test_mrtr_guards.py b/tests/server/test_mrtr_guards.py index 4855e2c5a..f5efa8309 100644 --- a/tests/server/test_mrtr_guards.py +++ b/tests/server/test_mrtr_guards.py @@ -22,6 +22,7 @@ from typing import Annotated import mcp_types import pytest +from docket import Docket from mcp.client._input_required import InputRequiredRoundsExceededError from mcp.server.request_state import RequestStateSecurity from mcp.shared.exceptions import MCPError @@ -37,6 +38,8 @@ from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware from fastmcp.server.middleware.middleware import Middleware from fastmcp.tools.base import InputRequiredToolResult, ToolResult from fastmcp.utilities.tests import run_server_async +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import running_task_server, submit_task, wait_for_task def _elicit(key: str, message: str, field: str) -> ElicitRequest: @@ -1158,9 +1161,18 @@ class TestTaskExecution: background task has no such request, so returning a guard result from a task is rejected with a clear error rather than silently yielding empty content.""" - @pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") - async def test_guard_result_from_task_is_rejected(self): + @pytest.fixture + def reset_docket_memory_server(self): + """Force a fresh memory:// Docket server bound to this test's loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + async def test_guard_result_from_task_is_rejected(self, reset_docket_memory_server): mcp = FastMCP("guard-task") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def book_flight(ctx: Context) -> str | InputRequiredResult: @@ -1170,13 +1182,14 @@ class TestTaskExecution: request_state=None, ) - # Client-side background-task submission (`task=True`) is the handshake-era - # SEP-1686 model; in 2026-07-28 tasks moved to a separate extension, so pin - # the era the "reject a guard's input-required from within a task" rule lives in. - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("book_flight", {}, task=True) + # A guard's `InputRequiredResult` only makes sense against a live + # request. Submitting `book_flight` as a background task and then + # reading it back must reject the guard result: `tasks/get` raises when + # it tries to inline the completed task's InputRequiredResult. + async with running_task_server(mcp): + created = await submit_task(mcp, "book_flight", {}) with pytest.raises(MCPError, match="background task"): - await task.result() + await wait_for_task(mcp, created.task_id) class TestHttpTransport: diff --git a/tests/server/test_server_docket.py b/tests/server/test_server_docket.py index ee11a5b60..bd4834b55 100644 --- a/tests/server/test_server_docket.py +++ b/tests/server/test_server_docket.py @@ -11,14 +11,21 @@ from fastmcp_tasks.dependencies import CurrentDocket, CurrentWorker from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.server.dependencies import get_context - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +from fastmcp_tasks import TasksExtension HUZZAH = "huzzah!" +@pytest.fixture(autouse=True) +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to each test's event loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + async def test_docket_not_initialized_without_task_components(): """Docket is only initialized when task-enabled components exist.""" mcp = FastMCP("test-server") @@ -28,10 +35,9 @@ async def test_docket_not_initialized_without_task_components(): return "no docket needed" async with Client(mcp) as client: - # Docket should not be initialized - assert mcp._docket is None + # Without a task=True tool, the lifespan never takes the Docket branch. + assert mcp.docket is None - # Regular tools still work result = await client.call_tool("regular_tool", {}) assert result.data == "no docket needed" @@ -39,8 +45,9 @@ async def test_docket_not_initialized_without_task_components(): async def test_current_docket(): """CurrentDocket dependency provides access to Docket instance.""" mcp = FastMCP("test-server") + mcp.add_extension(TasksExtension()) - # Need a task-enabled component to trigger Docket initialization + # A task-enabled component makes the lifespan start Docket. @mcp.tool(task=True) async def _trigger_docket() -> str: return "trigger" @@ -58,8 +65,8 @@ async def test_current_docket(): async def test_current_worker(): """CurrentWorker dependency provides access to Worker instance.""" mcp = FastMCP("test-server") + mcp.add_extension(TasksExtension()) - # Need a task-enabled component to trigger Docket initialization @mcp.tool(task=True) async def _trigger_docket() -> str: return "trigger" @@ -82,8 +89,8 @@ async def test_worker_executes_background_tasks(): """Verify that the Docket Worker is running and executes tasks.""" task_completed = asyncio.Event() mcp = FastMCP("test-server") + mcp.add_extension(TasksExtension()) - # Need a task-enabled component to trigger Docket initialization @mcp.tool(task=True) async def _trigger_docket() -> str: return "trigger" @@ -112,69 +119,12 @@ async def test_worker_executes_background_tasks(): await asyncio.wait_for(task_completed.wait(), timeout=2.0) -async def test_current_docket_in_resource(): - """CurrentDocket works in resources.""" - mcp = FastMCP("test-server") - - # Need a task-enabled component to trigger Docket initialization - @mcp.tool(task=True) - async def _trigger_docket() -> str: - return "trigger" - - @mcp.resource("docket://info") - def get_docket_info(docket: Docket = CurrentDocket()) -> str: - assert isinstance(docket, Docket) - return HUZZAH - - async with Client(mcp) as client: - result = await client.read_resource("docket://info") - assert HUZZAH in str(result) - - -async def test_current_docket_in_prompt(): - """CurrentDocket works in prompts.""" - mcp = FastMCP("test-server") - - # Need a task-enabled component to trigger Docket initialization - @mcp.tool(task=True) - async def _trigger_docket() -> str: - return "trigger" - - @mcp.prompt() - def task_prompt(task_type: str, docket: Docket = CurrentDocket()) -> str: - assert isinstance(docket, Docket) - return HUZZAH - - async with Client(mcp) as client: - result = await client.get_prompt("task_prompt", {"task_type": "background"}) - assert HUZZAH in str(result) - - -async def test_current_docket_in_resource_template(): - """CurrentDocket works in resource templates.""" - mcp = FastMCP("test-server") - - # Need a task-enabled component to trigger Docket initialization - @mcp.tool(task=True) - async def _trigger_docket() -> str: - return "trigger" - - @mcp.resource("docket://tasks/{task_id}") - def get_task_status(task_id: str, docket: Docket = CurrentDocket()) -> str: - assert isinstance(docket, Docket) - return HUZZAH - - async with Client(mcp) as client: - result = await client.read_resource("docket://tasks/123") - assert HUZZAH in str(result) - - async def test_concurrent_calls_maintain_isolation(): """Multiple concurrent calls each get the same Docket instance.""" mcp = FastMCP("test-server") + mcp.add_extension(TasksExtension()) docket_ids = [] - # Need a task-enabled component to trigger Docket initialization @mcp.tool(task=True) async def _trigger_docket() -> str: return "trigger" @@ -211,8 +161,8 @@ async def test_user_lifespan_still_works_with_docket(): yield {"custom_data": "test_value"} mcp = FastMCP("test-server", lifespan=custom_lifespan) + mcp.add_extension(TasksExtension()) - # Need a task-enabled component to trigger Docket initialization @mcp.tool(task=True) async def _trigger_docket() -> str: return "trigger" diff --git a/tests/server/test_tool_annotations.py b/tests/server/test_tool_annotations.py index 3c66040bd..44442fc76 100644 --- a/tests/server/test_tool_annotations.py +++ b/tests/server/test_tool_annotations.py @@ -1,11 +1,11 @@ from typing import Any -import pytest from mcp_types import Tool as MCPTool from mcp_types import ToolAnnotations, ToolExecution from fastmcp import Client, FastMCP from fastmcp.tools.base import Tool +from fastmcp_tasks import TasksExtension from tests.conftest import make_server_request_context @@ -221,25 +221,25 @@ async def test_tool_functionality_with_annotations(): assert result.data == {"name": "test_item", "value": 42} -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") async def test_task_execution_auto_populated_for_task_enabled_tool(): """Test that execution.task_support is automatically set when tool has task=True.""" mcp = FastMCP("Test Server") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def background_tool(data: str) -> str: """A tool that runs in background.""" return f"Processed: {data}" - # `execution.task_support` (SEP-1686) is advertised in the handshake-era - # tool listing only; the modern listing omits it. - async with Client(mcp, mode="legacy") as client: - tools_result = await client.list_tools() - assert len(tools_result) == 1 - assert tools_result[0].name == "background_tool" - assert isinstance(tools_result[0], MCPTool) - assert isinstance(tools_result[0].execution, ToolExecution) - assert tools_result[0].execution.task_support == "optional" + # The rendered tool descriptor auto-populates `execution.task_support` from + # the tool's task config. (The modern wire drops the SEP-1686 `execution` + # field, so this is asserted on the server-side render.) + tool = await mcp.get_tool("background_tool") + assert tool is not None + mcp_tool = tool.to_mcp_tool() + assert isinstance(mcp_tool, MCPTool) + assert isinstance(mcp_tool.execution, ToolExecution) + assert mcp_tool.execution.task_support == "optional" async def test_task_execution_omitted_for_task_disabled_tool(): diff --git a/tests/tasks/client/test_client_task_notifications.py b/tests/tasks/client/test_client_task_notifications.py index f93365caa..74cb5b207 100644 --- a/tests/tasks/client/test_client_task_notifications.py +++ b/tests/tasks/client/test_client_task_notifications.py @@ -16,9 +16,7 @@ from mcp_types import GetTaskResult from fastmcp import FastMCP from fastmcp.client import Client -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") async def _wait_until(condition: Callable[[], bool], timeout: float = 5.0) -> None: diff --git a/tests/tasks/client/test_client_task_protocol.py b/tests/tasks/client/test_client_task_protocol.py index 4d5d77bda..7b9698bb2 100644 --- a/tests/tasks/client/test_client_task_protocol.py +++ b/tests/tasks/client/test_client_task_protocol.py @@ -11,9 +11,7 @@ import pytest from fastmcp import FastMCP from fastmcp.client import Client -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") async def test_end_to_end_task_flow(): diff --git a/tests/tasks/client/test_client_tool_tasks.py b/tests/tasks/client/test_client_tool_tasks.py index 4a3d9d379..0a8220140 100644 --- a/tests/tasks/client/test_client_tool_tasks.py +++ b/tests/tasks/client/test_client_tool_tasks.py @@ -12,9 +12,7 @@ from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.exceptions import ToolError -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") @pytest.fixture diff --git a/tests/tasks/client/test_poll_interval.py b/tests/tasks/client/test_poll_interval.py index 0c33cca4d..fa4a25a4b 100644 --- a/tests/tasks/client/test_poll_interval.py +++ b/tests/tasks/client/test_poll_interval.py @@ -13,9 +13,7 @@ from fastmcp import Client, FastMCP from fastmcp.settings import Settings from fastmcp.utilities.tests import temporary_settings -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") @pytest.mark.parametrize("value", [0, -0.5, -1]) diff --git a/tests/tasks/client/test_task_context_validation.py b/tests/tasks/client/test_task_context_validation.py index 4eda41739..9d5f44e1e 100644 --- a/tests/tasks/client/test_task_context_validation.py +++ b/tests/tasks/client/test_task_context_validation.py @@ -10,9 +10,7 @@ import pytest from fastmcp import FastMCP from fastmcp.client import Client -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") @pytest.fixture diff --git a/tests/tasks/client/test_task_result_caching.py b/tests/tasks/client/test_task_result_caching.py index e0cf7b880..f7670cc6e 100644 --- a/tests/tasks/client/test_task_result_caching.py +++ b/tests/tasks/client/test_task_result_caching.py @@ -10,9 +10,7 @@ import pytest from fastmcp import FastMCP from fastmcp.client import Client -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +pytestmark = pytest.mark.skip(reason="Phase 4: requires client task support (SEP-2663)") async def test_tool_task_result_cached_on_first_call(): diff --git a/tests/tasks/server/conftest.py b/tests/tasks/server/conftest.py index 70ea6754b..b1dc7f98a 100644 --- a/tests/tasks/server/conftest.py +++ b/tests/tasks/server/conftest.py @@ -8,6 +8,26 @@ import pytest from fastmcp.utilities.tests import temporary_settings +@pytest.fixture(autouse=True) +def reset_docket_memory_server(): + """Reset the shared memory:// Docket server between tests. + + Docket keeps a process-wide ``Docket._memory_server`` singleton for + ``memory://`` backends. It persists across tests and across event loops, so a + test that inherits a stale server from a previous loop can fail (e.g. + ``tasks/get`` raising ``TypeError`` from the dead client). Clearing it before + and after each test keeps the task suite isolation-safe rather than + order-dependent. + """ + from docket import Docket + + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + @pytest.fixture(autouse=True) def isolate_settings_home(_settings_home_root: Path): """Task-local override of the repo-wide ``isolate_settings_home`` fixture. diff --git a/tests/tasks/server/test_concurrent_dependencies.py b/tests/tasks/server/test_concurrent_dependencies.py index 10db2860a..009723d25 100644 --- a/tests/tasks/server/test_concurrent_dependencies.py +++ b/tests/tasks/server/test_concurrent_dependencies.py @@ -8,37 +8,40 @@ Regression tests for: import asyncio -import pytest - from fastmcp import FastMCP -from fastmcp.client import Client from fastmcp.server.context import Context from fastmcp.server.dependencies import ( Progress, get_access_token, get_http_headers, ) - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + call_tool_without_optin, + running_task_server, + submit_task, + wait_for_task, ) async def test_concurrent_foreground_tools_with_context(): - """Multiple concurrent tool calls sharing the same CurrentContext() default + """Multiple concurrent tool calls sharing the same Context() default should not raise ValueError from ContextVar token resets (#3654).""" mcp = FastMCP("test") results: list[str] = [] - @mcp.tool() + @mcp.tool async def slow_tool(name: str, ctx: Context) -> str: await asyncio.sleep(0.01) results.append(name) return f"done:{name}" - async with Client(mcp, mode="legacy") as client: - tasks = [client.call_tool("slow_tool", {"name": f"task-{i}"}) for i in range(4)] - outcomes = await asyncio.gather(*tasks) + outcomes = await asyncio.gather( + *[ + call_tool_without_optin(mcp, "slow_tool", {"name": f"task-{i}"}) + for i in range(4) + ] + ) assert len(outcomes) == 4 for outcome in outcomes: @@ -50,7 +53,7 @@ async def test_concurrent_foreground_tools_with_progress(): should not raise AssertionError from _impl being None (#3656).""" mcp = FastMCP("test") - @mcp.tool() + @mcp.tool async def variable_tool( name: str, delay: float, progress: Progress = Progress() ) -> str: @@ -62,14 +65,14 @@ async def test_concurrent_foreground_tools_with_progress(): await progress.increment() return f"done:{name}" - async with Client(mcp, mode="legacy") as client: - tasks = [ - client.call_tool( - "variable_tool", {"name": f"t-{i}", "delay": 0.01 * (i + 1)} + outcomes = await asyncio.gather( + *[ + call_tool_without_optin( + mcp, "variable_tool", {"name": f"t-{i}", "delay": 0.01 * (i + 1)} ) for i in range(4) ] - outcomes = await asyncio.gather(*tasks) + ) assert len(outcomes) == 4 for outcome in outcomes: @@ -77,31 +80,33 @@ async def test_concurrent_foreground_tools_with_progress(): async def test_concurrent_background_tasks_with_context(): - """Multiple concurrent background tasks sharing _CurrentContext() should + """Multiple concurrent background tasks sharing Context() should not raise ValueError from ContextVar token resets (#3654).""" mcp = FastMCP("test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def bg_tool(name: str, ctx: Context) -> str: await asyncio.sleep(0.01) return f"bg:{name}" - async with Client(mcp, mode="legacy") as client: - task_handles = [ - await client.call_tool("bg_tool", {"name": f"bg-{i}"}, task=True) - for i in range(4) + async with running_task_server(mcp): + created = [ + await submit_task(mcp, "bg_tool", {"name": f"bg-{i}"}) for i in range(4) ] - results = await asyncio.gather(*[t.result() for t in task_handles]) + finals = await asyncio.gather(*[wait_for_task(mcp, c.task_id) for c in created]) - assert len(results) == 4 - for result in results: - assert result.content[0].text.startswith("bg:") + assert len(finals) == 4 + for final in finals: + assert final.status == "completed" + assert final.result["structuredContent"]["result"].startswith("bg:") async def test_concurrent_background_tasks_with_progress(): """Multiple concurrent background tasks sharing Progress() should not raise AssertionError from _impl being None (#3656).""" mcp = FastMCP("test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def bg_progress_tool( @@ -115,63 +120,62 @@ async def test_concurrent_background_tasks_with_progress(): await progress.increment() return f"bg:{name}" - async with Client(mcp, mode="legacy") as client: - task_handles = [ - await client.call_tool( + async with running_task_server(mcp): + created = [ + await submit_task( + mcp, "bg_progress_tool", {"name": f"bg-{i}", "delay": 0.01 * (i + 1)}, - task=True, ) for i in range(4) ] - results = await asyncio.gather(*[t.result() for t in task_handles]) + finals = await asyncio.gather(*[wait_for_task(mcp, c.task_id) for c in created]) - assert len(results) == 4 - for result in results: - assert result.content[0].text.startswith("bg:") + assert len(finals) == 4 + for final in finals: + assert final.status == "completed" + assert final.result["structuredContent"]["result"].startswith("bg:") async def test_dependency_aenter_returns_fresh_instances(): - """Verify that Dependency.__aenter__ returns independent per-invocation - objects, not the shared default.""" + """Dependency.__aenter__ returns independent per-invocation objects, + not the shared default.""" mcp = FastMCP("test") instances: list[Context] = [] - @mcp.tool() + @mcp.tool async def capture_context(ctx: Context) -> str: instances.append(ctx) return "ok" - async with Client(mcp, mode="legacy") as client: - await asyncio.gather( - client.call_tool("capture_context", {}), - client.call_tool("capture_context", {}), - ) + await asyncio.gather( + call_tool_without_optin(mcp, "capture_context", {}), + call_tool_without_optin(mcp, "capture_context", {}), + ) assert len(instances) == 2 assert instances[0] is not instances[1] async def test_progress_aenter_returns_fresh_instances(): - """Verify that Progress.__aenter__ returns independent per-invocation - objects, not the shared default.""" + """Progress.__aenter__ returns independent per-invocation objects, + not the shared default.""" progress_instances: list[Progress] = [] mcp = FastMCP("test") - @mcp.tool() + @mcp.tool async def capture_progress(progress: Progress = Progress()) -> str: progress_instances.append(progress) await progress.set_total(1) await progress.increment() return "ok" - async with Client(mcp, mode="legacy") as client: - await asyncio.gather( - client.call_tool("capture_progress", {}), - client.call_tool("capture_progress", {}), - ) + await asyncio.gather( + call_tool_without_optin(mcp, "capture_progress", {}), + call_tool_without_optin(mcp, "capture_progress", {}), + ) assert len(progress_instances) == 2 assert progress_instances[0] is not progress_instances[1] @@ -179,29 +183,32 @@ async def test_progress_aenter_returns_fresh_instances(): async def test_sync_context_functions_work_in_background_without_deps(): - """Sync functions like get_http_request() should work in background tasks - even when the tool declares no Context or CurrentRequest dependency. + """Sync helpers like get_http_headers() work in a background task even when + the tool declares no Context or CurrentRequest dependency. - This exercises the sync Redis fallback path (_get_task_snapshot_sync → - _load_snapshot_sync_redis) which must work with both memory:// (fakeredis) - and real Redis backends. + This exercises the sync snapshot fallback path which must work with the + memory:// (fakeredis) backend. """ mcp = FastMCP("test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def bare_sync_access() -> dict[str, str]: headers = get_http_headers() return {"has_headers": str(bool(headers))} - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("bare_sync_access", {}, task=True) - result = await task.result() - assert result.data == {"has_headers": "False"} + async with running_task_server(mcp): + created = await submit_task(mcp, "bare_sync_access", {}) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + assert final.result["structuredContent"] == {"has_headers": "False"} async def test_sync_context_functions_work_in_background_with_context(): - """Sync functions work via ContextVar when _CurrentContext loads the snapshot.""" + """Sync helpers work via ContextVar when Context loads the snapshot.""" mcp = FastMCP("test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def context_sync_access(ctx: Context) -> dict[str, str]: @@ -213,7 +220,9 @@ async def test_sync_context_functions_work_in_background_with_context(): "is_background": str(ctx.is_background_task), } - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("context_sync_access", {}, task=True) - result = await task.result() - assert result.data["is_background"] == "True" + async with running_task_server(mcp): + created = await submit_task(mcp, "context_sync_access", {}) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + assert final.result["structuredContent"]["is_background"] == "True" diff --git a/tests/tasks/server/test_context_background_task.py b/tests/tasks/server/test_context_background_task.py index 7f4bece2a..6cbc79e3f 100644 --- a/tests/tasks/server/test_context_background_task.py +++ b/tests/tasks/server/test_context_background_task.py @@ -1,63 +1,57 @@ -"""Tests for Context background task support (SEP-1686). +"""Tests for Context background task support (SEP-2663 tasks). -Tests Context API surface (unit) and background task elicitation (integration). -Integration tests use Client(mcp, mode="legacy") with the real memory:// Docket backend — -no mocking of Redis, Docket, or session internals. +Covers the Context API surface in a background task (unit tests, no Redis +needed) and end-to-end background-task behavior driven in-process through the +shared task helpers: progress reporting, context wiring, access-token +availability, and poll-based in-task elicitation. + +A SEP-2663 worker has no live session and no back-channel: ``ctx.session`` is +unavailable, and elicitation is polled (the worker parks an input request that +the client answers via ``tasks/update``). """ -import asyncio +from __future__ import annotations + import gc -import json from contextlib import AsyncExitStack -from datetime import datetime, timezone from typing import Any, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock import pytest -from fastmcp_tasks._legacy_wire.elicitation import handle_task_input from fastmcp_tasks.context import ( - TaskContextInfo, - TaskContextSnapshot, - _remember_snapshot, _task_sessions, - get_task_scope, get_task_session, register_task_session, ) -from fastmcp_tasks.dependencies import CurrentDocket -from fastmcp_tasks.keys import ( - task_redis_prefix, -) from mcp import ServerSession from mcp.server.auth.middleware.auth_context import auth_context_var from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from mcp_types import ( ClientCapabilities, - CreateMessageResult, Implementation, InitializeRequestParams, - TextContent, ) from pydantic import BaseModel from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.client.elicitation import ElicitResult from fastmcp.server.auth import AccessToken from fastmcp.server.context import Context from fastmcp.server.dependencies import get_access_token from fastmcp.server.elicitation import ( AcceptedElicitation, - CancelledElicitation, DeclinedElicitation, ) +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + running_task_server, + submit_task, + update_task, + wait_for_task, +) # ============================================================================= # Unit tests: Context API surface (no Redis/Docket needed) # ============================================================================= -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) class TestContextBackgroundTaskSupport: @@ -85,23 +79,9 @@ class TestContextBackgroundTaskSupport: setattr(ctx, "task_id", "new-id") -async def test_task_session_is_released_after_client_disconnect(): - _task_sessions.clear() - mcp = FastMCP("test") - - @mcp.tool(task=True) - async def work() -> str: - return "done" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("work", task=True) - await task.result() - assert len(_task_sessions) == 1 - - assert _task_sessions == {} - - async def test_live_task_session_is_released_on_connection_disconnect(): + """A registered in-process task session is dropped when its connection + exit stack unwinds.""" _task_sessions.clear() class MockConnection: @@ -124,6 +104,7 @@ async def test_live_task_session_is_released_on_connection_disconnect(): async def test_connection_cleanup_does_not_remove_replacement_session(): + """Registering a replacement session under the same id keeps the newer one.""" _task_sessions.clear() class MockConnection: @@ -147,6 +128,7 @@ async def test_connection_cleanup_does_not_remove_replacement_session(): def test_replaced_task_session_is_not_removed_by_old_weakref(): + """A stale weakref for a replaced session does not evict the new session.""" _task_sessions.clear() class MockSession: @@ -177,7 +159,7 @@ class TestContextSessionProperty: _ = ctx.session def test_session_uses_stored_session_in_background_task(self): - """session should use _session in background task mode.""" + """session should use the stored session in background task mode.""" mcp = FastMCP("test") class MockSession: @@ -191,7 +173,7 @@ class TestContextSessionProperty: 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).""" + """session should use the stored session during on_initialize.""" mcp = FastMCP("test") class MockSession: @@ -228,7 +210,7 @@ class TestContextBackgroundTaskLogging: return ctx, send_log_message async def test_background_task_honors_session_level(self): - """A background task has a session but no request context; the + """A background task has a stored session but no request context; the per-session minimum registered via logging/setLevel must still gate its logs, so sub-threshold messages are not sent to the client.""" mcp = FastMCP("test") @@ -258,10 +240,10 @@ class TestContextBackgroundTaskLogging: class TestContextClientExtensionBackgroundTask: """Tests for Context.client_supports_extension() in background task mode. - A background task has a live snapshot session but no request context. The - client's advertised capabilities are preserved on the snapshot session's - ``client_params``, so extension detection must read from the session rather - than gating on ``request_context``. + A background task may carry a stored snapshot session but no request + context. The client's advertised capabilities are preserved on the + session's ``client_params``, so extension detection reads from the session + rather than gating on ``request_context``. """ def _make_task_context( @@ -286,7 +268,7 @@ class TestContextClientExtensionBackgroundTask: ) def test_background_task_detects_advertised_extension(self): - """The snapshot session preserves the client's initialize params, so an + """The stored session preserves the client's initialize params, so an advertised extension is detected even with no request context.""" mcp = FastMCP("test") ctx = self._make_task_context(mcp, {"ext-abc": {}}) @@ -315,8 +297,9 @@ class TestContextClientExtensionBackgroundTask: class TestContextElicitBackgroundTask: """Tests for Context.elicit() in background task mode.""" - async def test_elicit_raises_when_background_task_but_no_docket(self): - """elicit() should raise when in background task mode but Docket unavailable.""" + async def test_elicit_raises_when_no_task_engine(self): + """elicit() fails fast when in a background task but no tasks extension + is installed to answer the request.""" mcp = FastMCP("test") ctx = Context(mcp, task_id="test-task-123") @@ -325,53 +308,10 @@ class TestContextElicitBackgroundTask: ctx._session = cast(ServerSession, MockSession()) - with pytest.raises(RuntimeError, match="Docket"): + with pytest.raises(RuntimeError, match="tasks extension"): 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. - """ - 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, mode="legacy") 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.""" @@ -392,143 +332,67 @@ class TestContextDocumentation: # ============================================================================= -# Integration tests: Client(mcp, mode="legacy") + memory:// Docket backend +# Integration tests: in-process SEP-2663 tasks via the shared helpers # ============================================================================= class TestBackgroundTaskIntegration: - """Integration tests for background task context using real Docket memory backend. - - These tests use Client(mcp, mode="legacy") with the memory:// broker — no mocking. - The memory:// backend provides a fully functional in-memory Redis store - that Docket uses automatically when running tests. - """ + """End-to-end background task context, driven in-process via the helpers.""" 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.add_extension(TasksExtension()) @mcp.tool(task=True) 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, mode="legacy") as client: - 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 with running_task_server(mcp): + created = await submit_task(mcp, "progress_tool", {}) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "done"} async def test_context_wiring_in_background_task(self): - """Context should be properly wired with task_id and session_id.""" + """A worker Context is wired as a background task with no live session.""" mcp = FastMCP("wiring-test") - task_completed = asyncio.Event() - captured: dict[str, object] = {} + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) - 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 "ok" + async def verify_wiring(ctx: Context) -> dict[str, bool]: + session_unavailable = False + try: + _ = ctx.session + except RuntimeError: + session_unavailable = True + return { + "task_id_set": ctx.task_id is not None, + "is_background": ctx.is_background_task, + "no_request_context": ctx.request_context is None, + "session_unavailable": session_unavailable, + } - async with Client(mcp, mode="legacy") as client: - 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 result.data == "ok" + async with running_task_server(mcp): + created = await submit_task(mcp, "verify_wiring", {}) + final = await wait_for_task(mcp, created.task_id) - assert captured["task_id"] is not None - assert captured["session_id"] is not None - assert captured["is_background"] is True - - async def test_origin_request_id_round_trips_through_background_task(self): - """E2E: origin_request_id captured at submit time is restored in worker. - - We validate this by comparing ctx.origin_request_id with the value - stored in Docket's Redis for this task. - """ - - mcp = FastMCP("origin-request-id-roundtrip") - - @mcp.tool(task=True) - async def check_origin_request_id(ctx: Context, docket=CurrentDocket()) -> str: - assert ctx.is_background_task is True - assert ctx.request_context is None - assert ctx.task_id is not None - - origin = ctx.origin_request_id - assert origin is not None - assert isinstance(origin, str) - assert origin != "" - - # Verify the snapshot in Redis contains the same value - task_scope = get_task_scope() - key = docket.key(f"{task_redis_prefix(task_scope)}:{ctx.task_id}:snapshot") - async with docket.redis() as redis: - raw = await redis.get(key) - - assert raw is not None - if isinstance(raw, bytes): - raw = raw.decode() - snapshot = json.loads(raw) - assert snapshot["origin_request_id"] == origin - return "ok" - - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("check_origin_request_id", {}, task=True) - result = await task.result() - assert result.data == "ok" - - @pytest.mark.xfail( - reason="Background-task sampling has no back-channel under SDK v2: the " - "per-request ServerSession that would carry sampling/createMessage is " - "gone once the submitting request completes, so ctx.sample() from a " - "worker raises NoBackChannelError. Needs a relay like elicit() " - "(context.py TODO); tracked in sdk-feedback.", - strict=True, - ) - async def test_sample_uses_origin_request_id_in_background_task(self): - """E2E: ctx.sample() works in a task without an active request context.""" - mcp = FastMCP("sample-background-test") - captured: dict[str, object] = {} - - @mcp.tool(task=True) - async def ask_client(ctx: Context) -> str: - assert ctx.is_background_task is True - assert ctx.request_context is None - assert ctx.origin_request_id is not None - result = await ctx.sample("Say hello") - return result.text or "" - - def sampling_handler(messages, params, ctx): - captured["called"] = True - return CreateMessageResult( - role="assistant", - content=TextContent(type="text", text="hello from background"), - model="test-model", - stop_reason="endTurn", - ) - - async with Client( - mcp, mode="legacy", sampling_handler=sampling_handler - ) as client: - task = await client.call_tool("ask_client", {}, task=True) - result = await task.result() - - assert result.data == "hello from background" - assert captured["called"] is True + assert final.status == "completed" + assert final.result["structuredContent"] == { + "task_id_set": True, + "is_background": True, + "no_request_context": True, + "session_unavailable": True, + } async def test_elicit_accept_flow(self): - """E2E: tool elicits input, client accepts via elicitation_handler.""" + """E2E: tool elicits input, client accepts via tasks/update (poll).""" mcp = FastMCP("elicit-accept-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def ask_name(ctx: Context) -> str: @@ -537,18 +401,26 @@ class TestBackgroundTaskIntegration: return f"Hello, {result.data}!" return "No name provided" - async def handler(message, response_type, params, ctx): - return ElicitResult(action="accept", content={"value": "Bob"}) + async with running_task_server(mcp): + created = await submit_task(mcp, "ask_name", {}) + parked = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"input_required"}) + ) + key = next(iter(parked.input_requests)) + await update_task( + mcp, + created.task_id, + {key: {"action": "accept", "content": {"value": "Bob"}}}, + ) + final = await wait_for_task(mcp, created.task_id) - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("ask_name", {}, task=True) - await task.wait(timeout=10.0) - result = await task.result() - assert result.data == "Hello, Bob!" + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "Hello, Bob!"} async def test_elicit_decline_flow(self): - """E2E: tool elicits input, client declines via elicitation_handler.""" + """E2E: tool elicits input, client declines via tasks/update (poll).""" mcp = FastMCP("elicit-decline-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def optional_input(ctx: Context) -> str: @@ -559,23 +431,27 @@ class TestBackgroundTaskIntegration: return f"Got: {result.data}" return "Cancelled" - async def handler(message, response_type, params, ctx): - return ElicitResult(action="decline") + async with running_task_server(mcp): + created = await submit_task(mcp, "optional_input", {}) + parked = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"input_required"}) + ) + key = next(iter(parked.input_requests)) + await update_task(mcp, created.task_id, {key: {"action": "decline"}}) + final = await wait_for_task(mcp, created.task_id) - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("optional_input", {}, task=True) - await task.wait(timeout=10.0) - result = await task.result() - assert result.data == "User declined" + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "User declined"} async def test_elicit_with_pydantic_model(self): - """E2E: tool elicits structured Pydantic input via elicitation_handler.""" + """E2E: tool elicits structured Pydantic input via tasks/update (poll).""" class UserInfo(BaseModel): name: str age: int mcp = FastMCP("elicit-pydantic-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def get_user_info(ctx: Context) -> str: @@ -585,51 +461,35 @@ class TestBackgroundTaskIntegration: return f"{result.data.name} is {result.data.age}" return "No info" - async def handler(message, response_type, params, ctx): - return ElicitResult(action="accept", content={"name": "Alice", "age": 30}) - - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("get_user_info", {}, task=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, mode="legacy") 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, - task_scope="nonexistent-scope", - action="accept", - content={"value": "too late"}, - fastmcp=mcp, + async with running_task_server(mcp): + created = await submit_task(mcp, "get_user_info", {}) + parked = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"input_required"}) ) - assert success is False + key = next(iter(parked.input_requests)) + await update_task( + mcp, + created.task_id, + {key: {"action": "accept", "content": {"name": "Alice", "age": 30}}}, + ) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "Alice is 30"} class TestAccessTokenInBackgroundTasks: """Tests for access token availability in background tasks (#3095). - Integration tests use Client(mcp, mode="legacy") with the real memory:// Docket backend. - The token snapshot/restore round-trip flows through actual Redis (fakeredis). - - Note: async tests run in isolated asyncio tasks, so ContextVar changes - are automatically scoped — no cleanup required. + The token set at submit time is available inside the worker (via the + captured context snapshot). Async tests run in isolated asyncio tasks, so + ContextVar changes are automatically scoped — no cleanup required. """ async def test_token_round_trips_through_background_task(self): """E2E: token set at submit time is available inside the worker.""" mcp = FastMCP("token-roundtrip") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def check_token(ctx: Context) -> str: @@ -646,81 +506,31 @@ class TestAccessTokenInBackgroundTasks: ) auth_context_var.set(AuthenticatedUser(test_token)) - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("check_token", {}, task=True) - result = await task.result() - assert result.data == "roundtrip-jwt|test-client" + async with running_task_server(mcp): + created = await submit_task(mcp, "check_token", {}) + final = await wait_for_task(mcp, created.task_id) + + assert final.status == "completed" + assert final.result["structuredContent"] == { + "result": "roundtrip-jwt|test-client" + } async def test_no_token_when_unauthenticated(self): """E2E: background task gets no token when nothing was set.""" mcp = FastMCP("no-auth") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def check_token(ctx: Context) -> str: token = get_access_token() return "no-token" if token is None else token.token - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("check_token", {}, task=True) - result = await task.result() - assert result.data == "no-token" + async with running_task_server(mcp): + created = await submit_task(mcp, "check_token", {}) + final = await wait_for_task(mcp, created.task_id) - async def test_expired_token_returns_none(self): - """get_access_token() returns None when task token has expired.""" - expired = AccessToken( - token="expired-jwt", - client_id="test-client", - scopes=["read"], - expires_at=int(datetime.now(timezone.utc).timestamp()) - 3600, - ) - _remember_snapshot( - "test-task", - TaskContextSnapshot(access_token_json=expired.model_dump_json()), - ) - fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s") - with patch( - "fastmcp.server.dependencies.get_task_context", return_value=fake_ctx - ): - assert get_access_token() is None - - async def test_valid_token_with_future_expiry(self): - """get_access_token() returns token when expiry is in the future.""" - valid = AccessToken( - token="valid-jwt", - client_id="test-client", - scopes=["read"], - expires_at=int(datetime.now(timezone.utc).timestamp()) + 3600, - ) - _remember_snapshot( - "test-task", - TaskContextSnapshot(access_token_json=valid.model_dump_json()), - ) - fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s") - with patch( - "fastmcp.server.dependencies.get_task_context", return_value=fake_ctx - ): - result = get_access_token() - assert result is not None - assert result.token == "valid-jwt" - - async def test_token_without_expiry_always_valid(self): - """get_access_token() returns token when no expires_at is set.""" - no_expiry = AccessToken( - token="eternal-jwt", - client_id="test-client", - scopes=["read"], - ) - _remember_snapshot( - "test-task", - TaskContextSnapshot(access_token_json=no_expiry.model_dump_json()), - ) - fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s") - with patch( - "fastmcp.server.dependencies.get_task_context", return_value=fake_ctx - ): - result = get_access_token() - assert result is not None - assert result.token == "eternal-jwt" + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "no-token"} class TestLifespanContextInBackgroundTasks: diff --git a/tests/tasks/server/test_custom_subclass_tasks.py b/tests/tasks/server/test_custom_subclass_tasks.py index cd6e87a39..c9b41fafc 100644 --- a/tests/tasks/server/test_custom_subclass_tasks.py +++ b/tests/tasks/server/test_custom_subclass_tasks.py @@ -1,7 +1,8 @@ -"""Tests for custom component subclasses with task support. +"""Tests for custom Tool subclasses with task support. -Verifies that custom Tool, Resource, and Prompt subclasses can use -background task execution by setting task_config. +Verifies that custom Tool subclasses can use background task execution by +setting task_config. SEP-2663 is tools-only, so the removed resource/prompt +subclass cases are gone. """ import asyncio @@ -9,15 +10,23 @@ from typing import Any from unittest.mock import MagicMock import pytest +from fastmcp_tasks.components import ( + add_component_to_docket, + register_component_with_docket, +) +from fastmcp_tasks.models import CreateTaskResult from fastmcp import FastMCP -from fastmcp.client import Client from fastmcp.tools.base import Tool, ToolResult from fastmcp.utilities.components import FastMCPComponent from fastmcp.utilities.tasks import TaskConfig - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + _opted_in_request, + auth_scope, + call_tool_without_optin, + run_task, + running_task_server, ) @@ -56,9 +65,10 @@ class CustomToolForbidden(Tool): @pytest.fixture -def custom_tool_server(): - """Create a server with custom tool subclasses.""" +def custom_tool_server() -> FastMCP: + """A server with custom tool subclasses.""" mcp = FastMCP("custom-tool-server") + mcp.add_extension(TasksExtension()) mcp.add_tool(CustomTool(name="custom_tool", description="A custom tool")) mcp.add_tool( CustomToolWithLogic(name="custom_logic", description="Custom tool with logic") @@ -70,76 +80,67 @@ def custom_tool_server(): async def test_custom_tool_sync_execution(custom_tool_server): - """Custom tool executes synchronously when no task metadata.""" - async with Client(custom_tool_server, mode="legacy") as client: - result = await client.call_tool("custom_tool", {}) - assert "Custom tool executed" in str(result) + """Custom tool executes synchronously without a tasks opt-in.""" + async with running_task_server(custom_tool_server): + result = await call_tool_without_optin(custom_tool_server, "custom_tool", {}) + assert "Custom tool executed" in result.content[0].text async def test_custom_tool_background_execution(custom_tool_server): - """Custom tool executes as background task when task=True.""" - async with Client(custom_tool_server, mode="legacy") as client: - task = await client.call_tool("custom_tool", {}, task=True) + """Custom tool executes as a background task when opted in.""" + async with running_task_server(custom_tool_server): + final = await run_task(custom_tool_server, "custom_tool", {}) - assert task is not None - assert not task.returned_immediately - assert task.task_id is not None - - # Wait for result - result = await task.result() - assert "Custom tool executed" in str(result) + assert final.status == "completed" + assert "Custom tool executed" in final.result["content"][0]["text"] async def test_custom_tool_with_arguments(custom_tool_server): """Custom tool receives arguments correctly in background execution.""" - async with Client(custom_tool_server, mode="legacy") as client: - task = await client.call_tool("custom_logic", {"duration": 1}, task=True) + async with running_task_server(custom_tool_server): + final = await run_task(custom_tool_server, "custom_logic", {"duration": 1}) - assert task is not None - result = await task.result() - assert "Completed after 1 units" in str(result) + assert final.status == "completed" + assert "Completed after 1 units" in final.result["content"][0]["text"] async def test_custom_tool_forbidden_sync_only(custom_tool_server): - """Custom tool with forbidden mode executes sync only.""" - async with Client(custom_tool_server, mode="legacy") as client: - # Sync execution works - result = await client.call_tool("custom_forbidden", {}) - assert "Sync only" in str(result) + """Custom tool with forbidden mode executes synchronously.""" + async with running_task_server(custom_tool_server): + result = await call_tool_without_optin( + custom_tool_server, "custom_forbidden", {} + ) + assert "Sync only" in result.content[0].text async def test_custom_tool_forbidden_rejects_task(custom_tool_server): - """Custom tool with forbidden mode returns error for task request.""" - async with Client(custom_tool_server, mode="legacy") as client: - task = await client.call_tool("custom_forbidden", {}, task=True) - - # Should return immediately with error - assert task.returned_immediately + """A forbidden tool runs synchronously even when the client opts in.""" + async with running_task_server(custom_tool_server): + with auth_scope(None), _opted_in_request("custom_forbidden", {}, None): + result = await custom_tool_server.call_tool("custom_forbidden", {}) + assert not isinstance(result, CreateTaskResult) + assert "Sync only" in result.content[0].text async def test_custom_tool_registers_with_docket(): - """Verify custom tool's register_with_docket is called during server startup.""" - from unittest.mock import MagicMock - + """A task-capable custom tool registers its `run` entry point with Docket.""" tool = CustomTool(name="test", description="test") mock_docket = MagicMock() - tool.register_with_docket(mock_docket) + register_component_with_docket(tool, mock_docket) - # Should register self.run with docket using prefixed key mock_docket.register.assert_called_once() call_args = mock_docket.register.call_args assert call_args[1]["names"] == ["tool:test@"] async def test_custom_tool_forbidden_does_not_register(): - """Verify custom tool with forbidden mode doesn't register with docket.""" + """A forbidden custom tool does not register with Docket.""" tool = CustomToolForbidden(name="test", description="test") mock_docket = MagicMock() - tool.register_with_docket(mock_docket) + register_component_with_docket(tool, mock_docket) - # Should NOT register mock_docket.register.assert_not_called() @@ -157,26 +158,24 @@ class TestFastMCPComponentDocketMethods: assert component.task_config.mode == "forbidden" def test_register_with_docket_is_noop(self): - """Base register_with_docket does nothing (subclasses override).""" + """Registering a forbidden base component is a no-op.""" component = FastMCPComponent(name="test") mock_docket = MagicMock() - # Should not raise, just no-op - component.register_with_docket(mock_docket) + register_component_with_docket(component, mock_docket) - # Should not have called any docket methods mock_docket.register.assert_not_called() async def test_add_to_docket_raises_when_forbidden(self): - """Base add_to_docket raises RuntimeError when mode is 'forbidden'.""" + """add_component_to_docket raises RuntimeError when mode is 'forbidden'.""" component = FastMCPComponent(name="test") mock_docket = MagicMock() with pytest.raises(RuntimeError, match="task execution not supported"): - await component.add_to_docket(mock_docket) + await add_component_to_docket(component, mock_docket, None) async def test_add_to_docket_raises_not_implemented_when_allowed(self): - """Base add_to_docket raises NotImplementedError when not forbidden.""" + """add_component_to_docket raises NotImplementedError for an unknown type.""" component = FastMCPComponent( name="test", task_config=TaskConfig(mode="optional") ) @@ -185,4 +184,4 @@ class TestFastMCPComponentDocketMethods: with pytest.raises( NotImplementedError, match="does not implement add_to_docket" ): - await component.add_to_docket(mock_docket) + await add_component_to_docket(component, mock_docket, None) diff --git a/tests/tasks/server/test_extension.py b/tests/tasks/server/test_extension.py new file mode 100644 index 000000000..64eb9ad5a --- /dev/null +++ b/tests/tasks/server/test_extension.py @@ -0,0 +1,366 @@ +"""End-to-end tests for the SEP-2663 `TasksExtension` server adapter. + +Covers the decide-and-task interceptor (forbidden/optional/required modes and the +-32003 missing-capability error), the tasks/get|update|cancel handlers, status +mapping, inlined completed results, argument-coercion parity, TTL, and capability +advertisement. Server-side tasks are driven in-process via `task_helpers` because +there is no client task-submission API until Phase 4. +""" + +from __future__ import annotations + +import asyncio +from contextlib import AsyncExitStack +from types import SimpleNamespace + +import pytest +from fastmcp_tasks.models import ( + MISSING_REQUIRED_CLIENT_CAPABILITY, + CreateTaskResult, +) +from mcp.server.context import ServerRequestContext +from mcp.shared.exceptions import MCPError + +from fastmcp import FastMCP +from fastmcp.client import Client +from fastmcp.exceptions import ToolError +from fastmcp.server import context as core_context +from fastmcp.server.dependencies import bind_request_context +from fastmcp.tools.base import ToolResult +from fastmcp.utilities.tasks import TASKS_EXTENSION_ID, TaskConfig +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + _opted_in_request, + auth_scope, + call_tool_without_optin, + get_task, + make_access_token, + opt_in_meta, + run_task, + running_task_server, + submit_task, + wait_for_task, +) + + +def _tasks_server() -> FastMCP: + mcp = FastMCP("tasks") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def square(n: int) -> int: + return n * n + + @mcp.tool(task=TaskConfig(mode="required")) + async def must_task(n: int) -> int: + return n + 1 + + @mcp.tool + async def plain(n: int) -> int: + return n - 1 + + @mcp.tool(task=True) + async def boom() -> int: + raise ToolError("kaboom") + + return mcp + + +# --------------------------------------------------------------------------- +# Capability advertisement +# --------------------------------------------------------------------------- + + +async def test_capability_advertised_to_modern_client(): + mcp = FastMCP("t") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def t(n: int) -> int: + return n + + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert extensions.get(TASKS_EXTENSION_ID) == {} + + +async def test_capability_absent_without_extension(): + mcp = FastMCP("t") + + @mcp.tool + async def t(n: int) -> int: + return n + + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert TASKS_EXTENSION_ID not in extensions + + +# --------------------------------------------------------------------------- +# Decide-and-task interceptor +# --------------------------------------------------------------------------- + + +async def test_optional_tool_tasks_when_opted_in(): + mcp = _tasks_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "square", {"n": 5}) + assert isinstance(created, CreateTaskResult) + assert created.status == "working" + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result is not None + assert final.result["structuredContent"]["result"] == 25 + + +async def test_optional_tool_runs_sync_without_opt_in(): + mcp = _tasks_server() + async with running_task_server(mcp): + result = await call_tool_without_optin(mcp, "square", {"n": 5}) + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": 25} + + +async def test_forbidden_tool_never_tasks_even_with_opt_in(): + mcp = _tasks_server() + async with running_task_server(mcp): + # `plain` is mode=forbidden; opting in must not task it. + result = await submit_task_expecting_sync(mcp, "plain", {"n": 5}) + assert result.structured_content == {"result": 4} + + +async def submit_task_expecting_sync(mcp, name, args): + with auth_scope(None), _opted_in_request(name, args, None): + return await mcp.call_tool(name, args) + + +async def test_required_tool_tasks_when_opted_in(): + mcp = _tasks_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "must_task", {"n": 10}) + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result["structuredContent"]["result"] == 11 + + +async def test_required_tool_without_opt_in_raises_missing_capability(): + mcp = _tasks_server() + async with running_task_server(mcp): + with pytest.raises(MCPError) as exc_info: + await call_tool_without_optin(mcp, "must_task", {"n": 1}) + error = exc_info.value.error + assert error.code == MISSING_REQUIRED_CLIENT_CAPABILITY + assert error.data == { + "requiredCapabilities": {"extensions": {TASKS_EXTENSION_ID: {}}} + } + + +# --------------------------------------------------------------------------- +# Task id and status +# --------------------------------------------------------------------------- + + +async def test_task_ids_are_server_generated_and_distinct(): + mcp = _tasks_server() + async with running_task_server(mcp): + a = await submit_task(mcp, "square", {"n": 1}) + b = await submit_task(mcp, "square", {"n": 2}) + assert a.task_id != b.task_id + assert len(a.task_id) >= 20 + + +async def test_get_unknown_task_raises_not_found(): + mcp = _tasks_server() + async with running_task_server(mcp): + with pytest.raises(MCPError, match="not found"): + await get_task(mcp, "does-not-exist") + + +async def test_failed_task_surfaces_error_not_completed(): + mcp = _tasks_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "boom", {}) + final = await wait_for_task(mcp, created.task_id) + assert final.status == "failed" + assert final.error is not None + assert "kaboom" in final.error["message"] + assert final.result is None + + +# --------------------------------------------------------------------------- +# Argument coercion parity +# --------------------------------------------------------------------------- + + +async def test_task_arguments_are_coerced_like_sync_path(): + mcp = _tasks_server() + async with running_task_server(mcp): + # "6" coerces to int 6 exactly as the synchronous path would. + final = await run_task(mcp, "square", {"n": "6"}) + assert final.result["structuredContent"]["result"] == 36 + + +async def test_invalid_task_arguments_reject_at_submission(): + mcp = _tasks_server() + async with running_task_server(mcp): + with pytest.raises(Exception): + await submit_task(mcp, "square", {"n": "not-a-number"}) + + +# --------------------------------------------------------------------------- +# TTL / poll interval +# --------------------------------------------------------------------------- + + +async def test_create_and_get_carry_ttl_and_poll_interval(): + mcp = _tasks_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "square", {"n": 3}) + assert created.ttl_ms is not None and created.ttl_ms > 0 + assert created.poll_interval_ms == 5000 + got = await get_task(mcp, created.task_id) + assert got.ttl_ms == created.ttl_ms + assert got.poll_interval_ms == 5000 + + +# --------------------------------------------------------------------------- +# Cancellation +# --------------------------------------------------------------------------- + + +async def test_cancel_transitions_task_to_cancelled(): + mcp = FastMCP("t") + mcp.add_extension(TasksExtension()) + release = asyncio.Event() + + @mcp.tool(task=True) + async def slow() -> str: + await release.wait() + return "done" + + async with running_task_server(mcp): + created = await submit_task(mcp, "slow", {}) + ack = await cancel_and_release(mcp, created.task_id, release) + assert ack is not None + final = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"cancelled", "completed"}) + ) + assert final.status in {"cancelled", "completed"} + + +async def cancel_and_release(mcp, task_id, release): + from tests.tasks.task_helpers import cancel_task + + ack = await cancel_task(mcp, task_id) + release.set() + return ack + + +# --------------------------------------------------------------------------- +# Serve-time guard +# --------------------------------------------------------------------------- + + +async def test_task_tool_without_extension_fails_at_serve_time(): + mcp = FastMCP("t") + + @mcp.tool(task=True) + async def t(n: int) -> int: + return n + + with pytest.raises(RuntimeError, match="tasks extension"): + async with mcp._lifespan_manager(): + pass + + +# --------------------------------------------------------------------------- +# Auth-scoped isolation +# --------------------------------------------------------------------------- + + +async def test_tasks_isolated_across_auth_scopes(): + mcp = _tasks_server() + alice = make_access_token("alice") + bob = make_access_token("bob") + async with running_task_server(mcp): + created = await submit_task(mcp, "square", {"n": 4}, access_token=alice) + # Alice sees her task. + mine = await get_task(mcp, created.task_id, access_token=alice) + assert mine.task_id == created.task_id + # Bob cannot: a cross-scope id is indistinguishable from missing. + with pytest.raises(MCPError, match="not found"): + await get_task(mcp, created.task_id, access_token=bob) + + +# --------------------------------------------------------------------------- +# Protocol-era gating of the tasking decision +# --------------------------------------------------------------------------- + + +async def test_legacy_era_opt_in_is_ignored(): + """A handshake-era request cannot be tasked, even with the _meta opt-in. + + The SDK strips `capabilities.extensions` from pre-2026 handshakes, so a + legacy client can never have negotiated the tasks extension — a stray + per-request opt-in on a legacy connection is treated as absent and an + `optional` tool runs synchronously. + """ + mcp = _tasks_server() + async with running_task_server(mcp): + srctx = ServerRequestContext( + session=SimpleNamespace(), + lifespan_context={}, + protocol_version="2025-06-18", + method="tools/call", + params={"name": "square", "arguments": {"n": 3}, "_meta": opt_in_meta()}, + ) + with bind_request_context(srctx): + result = await mcp.call_tool("square", {"n": 3}) + assert isinstance(result, ToolResult) + + +async def test_legacy_era_required_tool_raises_missing_capability(): + """`required` tools refuse legacy-era calls with -32003 even when opted in.""" + mcp = _tasks_server() + async with running_task_server(mcp): + srctx = ServerRequestContext( + session=SimpleNamespace(), + lifespan_context={}, + protocol_version="2025-06-18", + method="tools/call", + params={ + "name": "must_task", + "arguments": {"n": 3}, + "_meta": opt_in_meta(), + }, + ) + with bind_request_context(srctx): + with pytest.raises(MCPError) as exc_info: + await mcp.call_tool("must_task", {"n": 3}) + assert exc_info.value.error.code == -32003 + + +# --------------------------------------------------------------------------- +# Worker-hook lifecycle across multiple servers +# --------------------------------------------------------------------------- + + +async def test_worker_hooks_survive_sibling_server_shutdown(): + """One server's shutdown must not strand another server's workers. + + The worker-side hooks core exposes are process-global; two sibling servers + each running a TasksExtension refcount them, so the hooks clear only when + the last extension lifespan exits. + """ + server_a = _tasks_server() + server_b = _tasks_server() + + async with AsyncExitStack() as stack_b: + await stack_b.enter_async_context(server_b._lifespan_manager()) + async with AsyncExitStack() as stack_a: + await stack_a.enter_async_context(server_a._lifespan_manager()) + assert core_context._task_elicitation_handler is not None + # Server A has shut down; server B's workers still need the hooks. + assert core_context._task_elicitation_handler is not None + # The last extension exited; hooks are cleared. + assert core_context._task_elicitation_handler is None diff --git a/tests/tasks/server/test_notifications.py b/tests/tasks/server/test_notifications.py deleted file mode 100644 index 32c6c90bc..000000000 --- a/tests/tasks/server/test_notifications.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Tests for distributed notification queue (SEP-1686). - -Integration tests verify that the notification queue works end-to-end -using Client(mcp, mode="legacy") with the real memory:// Docket backend. -No mocking of Redis, sessions, or Docket internals. -""" - -import asyncio -import time - -import mcp_types -import pytest -from fastmcp_tasks._legacy_wire.notifications import ( - get_subscriber_count, -) - -from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.client.elicitation import ElicitResult -from fastmcp.server.context import Context -from fastmcp.server.elicitation import AcceptedElicitation - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) - - -class TestNotificationIntegration: - """Integration tests for the notification queue using real Docket memory backend. - - The elicitation flow 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. Subscriber relays elicitation/create to client -> handler responds - 4. Relay pushes response to Redis -> BLPOP wakes tool - """ - - async def test_notification_delivered_during_elicitation(self): - """Full E2E: notification queue delivers input_required metadata to client. - - SDK v2 does not carry `notifications/tasks/status` in any protocol - version's core notification tables, so it is delivered through the - client's task-status notification binding (routed to Task objects) rather - than the message_handler. We observe it via `on_status_change`, whose - GetTaskResult carries the notification's `_meta`. - """ - mcp = FastMCP("notification-test") - captured: list[mcp_types.GetTaskResult] = [] - - @mcp.tool(task=True) - async def elicit_tool(ctx: Context) -> str: - result = await ctx.elicit("Enter value", str) - if isinstance(result, AcceptedElicitation): - return f"got: {result.data}" - return "no value" - - async def elicitation_handler(message, response_type, params, ctx): - return ElicitResult(action="accept", content={"value": "hello"}) - - async with Client( - mcp, - mode="legacy", - elicitation_handler=elicitation_handler, - ) as client: - task = await client.call_tool("elicit_tool", {}, task=True) - task.on_status_change(captured.append) - - await task.wait(timeout=10.0) - result = await task.result() - assert result.data == "got: hello" - - # Verify the input_required notification was delivered with metadata - notification: mcp_types.GetTaskResult | None = None - for candidate in reversed(captured): - candidate_meta = candidate.meta - related_task = ( - candidate_meta.get("io.modelcontextprotocol/related-task") - if isinstance(candidate_meta, dict) - else None - ) - if ( - isinstance(related_task, dict) - and related_task.get("status") == "input_required" - ): - notification = candidate - break - - assert notification is not None, "expected notifications/tasks/status" - task_meta = notification.meta - assert isinstance(task_meta, dict) - - related_task = task_meta.get("io.modelcontextprotocol/related-task") - assert isinstance(related_task, dict) - assert related_task.get("taskId") == task.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) - - 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() - 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, mode="legacy") 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 - deadline = time.monotonic() + 1.0 - while get_subscriber_count() != count_before and time.monotonic() < deadline: - await asyncio.sleep(0.005) - assert get_subscriber_count() == count_before diff --git a/tests/tasks/server/test_progress_dependency.py b/tests/tasks/server/test_progress_dependency.py index 3cf751eb0..98401d9d2 100644 --- a/tests/tasks/server/test_progress_dependency.py +++ b/tests/tasks/server/test_progress_dependency.py @@ -1,38 +1,41 @@ -"""Tests for FastMCP Progress dependency.""" +"""Tests for FastMCP Progress dependency (SEP-2663 tasks).""" -import pytest +import asyncio +import json + +from mcp_types import TextContent from fastmcp import FastMCP -from fastmcp.client import Client from fastmcp.server.dependencies import Progress - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + call_tool_without_optin, + running_task_server, + submit_task, + wait_for_task, ) async def test_progress_in_immediate_execution(): - """Test Progress dependency when calling tool immediately with Docket enabled.""" + """Progress dependency works when a tool runs synchronously.""" mcp = FastMCP("test") - @mcp.tool() + @mcp.tool async def test_tool(progress: Progress = Progress()) -> str: await progress.set_total(10) await progress.increment() await progress.set_message("Testing") return "done" - async with Client(mcp, mode="legacy") as client: - result = await client.call_tool("test_tool", {}) - from mcp_types import TextContent - - assert isinstance(result.content[0], TextContent) - assert result.content[0].text == "done" + result = await call_tool_without_optin(mcp, "test_tool", {}) + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "done" async def test_progress_in_background_task(): - """Test Progress dependency in background task execution.""" + """Progress dependency works inside a background task.""" mcp = FastMCP("test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def test_task(progress: Progress = Progress()) -> str: @@ -41,104 +44,81 @@ async def test_progress_in_background_task(): await progress.set_message("Step 1") return "done" - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("test_task", {}, task=True) - result = await task.result() - from mcp_types import TextContent - - assert isinstance(result.content[0], TextContent) - assert result.content[0].text == "done" + async with running_task_server(mcp): + created = await submit_task(mcp, "test_task", {}) + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "done"} async def test_progress_tracks_multiple_increments(): - """Test that Progress correctly tracks multiple increment calls.""" + """Progress correctly tracks multiple increment calls.""" mcp = FastMCP("test") - @mcp.tool() + @mcp.tool async def count_to_ten(progress: Progress = Progress()) -> str: await progress.set_total(10) - for i in range(10): + for _ in range(10): await progress.increment() return "counted" - async with Client(mcp, mode="legacy") as client: - result = await client.call_tool("count_to_ten", {}) - from mcp_types import TextContent - - assert isinstance(result.content[0], TextContent) - assert result.content[0].text == "counted" + result = await call_tool_without_optin(mcp, "count_to_ten", {}) + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "counted" async def test_progress_status_message_in_background_task(): - """Regression test: TaskStatusResponse must include statusMessage field.""" - import asyncio - + """A working task surfaces the current progress message as statusMessage.""" mcp = FastMCP("test") - step_started = asyncio.Event() + mcp.add_extension(TasksExtension()) + release = asyncio.Event() @mcp.tool(task=True) async def task_with_progress(progress: Progress = Progress()) -> str: await progress.set_total(3) await progress.set_message("Step 1 of 3") await progress.increment() - step_started.set() - - # No settling wait needed: the server never clears the progress - # message on completion (only a failure overwrites it), so whatever - # "Step N of 3" message is current when the test polls status() - # below still satisfies the assertion, win or lose the race. + await release.wait() await progress.set_message("Step 2 of 3") await progress.increment() - await progress.set_message("Step 3 of 3") - await progress.increment() return "done" - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("task_with_progress", {}, task=True) + async with running_task_server(mcp): + created = await submit_task(mcp, "task_with_progress", {}) - # Wait for first step to start - await step_started.wait() - - # Get status and verify progress message - status = await task.status() - - # Verify statusMessage field is accessible and contains progress info - # Should not raise AttributeError - msg = status.status_message + # The task parks on `release` while working; its statusMessage should + # reflect the progress message (or be None, depending on the poll race). + working = await wait_for_task( + mcp, created.task_id, target_states=frozenset({"working"}) + ) + msg = working.status_message assert msg is None or msg.startswith("Step") - # Wait for completion - result = await task.result() - from mcp_types import TextContent - - assert isinstance(result.content[0], TextContent) - assert result.content[0].text == "done" + release.set() + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "done"} async def test_inmemory_progress_state(): - """Test that in-memory progress stores and returns state correctly.""" + """In-memory progress stores and returns state correctly.""" mcp = FastMCP("test") - @mcp.tool() + @mcp.tool async def test_tool(progress: Progress = Progress()) -> dict: - # Initial state assert progress.current is None assert progress.total == 1 assert progress.message is None - # Set total await progress.set_total(10) assert progress.total == 10 - # Increment await progress.increment() assert progress.current == 1 - # Increment again await progress.increment(2) assert progress.current == 3 - # Set message await progress.set_message("Testing") assert progress.message == "Testing" @@ -148,15 +128,9 @@ async def test_inmemory_progress_state(): "message": progress.message, } - async with Client(mcp, mode="legacy") as client: - result = await client.call_tool("test_tool", {}) - from mcp_types import TextContent - - assert isinstance(result.content[0], TextContent) - # The tool returns a dict showing the final state - import json - - state = json.loads(result.content[0].text) - assert state["current"] == 3 - assert state["total"] == 10 - assert state["message"] == "Testing" + result = await call_tool_without_optin(mcp, "test_tool", {}) + assert isinstance(result.content[0], TextContent) + state = json.loads(result.content[0].text) + assert state["current"] == 3 + assert state["total"] == 10 + assert state["message"] == "Testing" diff --git a/tests/tasks/server/test_server_tasks_parameter.py b/tests/tasks/server/test_server_tasks_parameter.py index 3f79a4820..6ebb6d121 100644 --- a/tests/tasks/server/test_server_tasks_parameter.py +++ b/tests/tasks/server/test_server_tasks_parameter.py @@ -1,445 +1,135 @@ -""" -Tests for server `tasks` parameter default inheritance. +"""Server-level `tasks` default inheritance and per-tool override (tools only). -Verifies that the server's `tasks` parameter correctly sets defaults for all -components (tools, prompts, resources), and that explicit component-level -settings properly override the server default. +`FastMCP(tasks=...)` sets the default task mode for tools; a per-tool `task=` +overrides it. SEP-2663 tasks are tools-only, so prompt/resource/template +inheritance is not covered. Tasking is driven in-process through the interceptor. """ -import pytest +from __future__ import annotations + +from fastmcp_tasks.models import CreateTaskResult from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + _opted_in_request, + auth_scope, + run_task, + running_task_server, + submit_task, ) -@pytest.mark.timeout(10) -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_server_tasks_true_defaults_all_components(): - """Server with tasks=True makes all components default to supporting tasks.""" +async def _opted_in_call(server: FastMCP, name: str, arguments: dict | None = None): + """Run a `tools/call` WITH the tasks opt-in bound (used to prove sync paths).""" + with auth_scope(None), _opted_in_request(name, arguments or {}, None): + return await server.call_tool(name, arguments or {}) + + +async def test_tool_inherits_server_default_true(): + """A tool inherits the server's tasks=True default and tasks when opted in.""" mcp = FastMCP("test", tasks=True) + mcp.add_extension(TasksExtension()) - @mcp.tool() + @mcp.tool async def my_tool() -> str: return "tool result" - @mcp.prompt() - async def my_prompt() -> str: - return "prompt result" - - @mcp.resource("test://resource") - async def my_resource() -> str: - return "resource result" - - async with Client(mcp, mode="legacy") as client: - # Verify all task-enabled components are registered with docket - # Components use prefixed keys: tool:name, prompt:name, resource:uri - docket = mcp.docket - assert docket is not None - assert "tool:my_tool@" in docket.tasks - assert "prompt:my_prompt@" in docket.tasks - assert "resource:test://resource@" in docket.tasks - - # Tool should support background execution - tool_task = await client.call_tool("my_tool", task=True) - assert not tool_task.returned_immediately - - # Prompt should support background execution - prompt_task = await client.get_prompt("my_prompt", task=True) - assert not prompt_task.returned_immediately - - # Resource should support background execution - resource_task = await client.read_resource("test://resource", task=True) - assert not resource_task.returned_immediately + async with running_task_server(mcp): + created = await submit_task(mcp, "my_tool") + assert isinstance(created, CreateTaskResult) -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_server_tasks_false_defaults_all_components(): - """Server with tasks=False makes all components default to mode=forbidden.""" - import pytest - from mcp.shared.exceptions import MCPError - +async def test_tool_inherits_server_default_false(): + """A tool inherits the server's tasks=False default and runs synchronously.""" mcp = FastMCP("test", tasks=False) - @mcp.tool() + @mcp.tool async def my_tool() -> str: return "tool result" - @mcp.prompt() - async def my_prompt() -> str: - return "prompt result" - - @mcp.resource("test://resource") - async def my_resource() -> str: - return "resource result" - - async with Client(mcp, mode="legacy") as client: - # Tool with mode="forbidden" returns error when called with task=True - tool_task = await client.call_tool("my_tool", task=True, raise_on_error=False) - assert tool_task.returned_immediately - result = await tool_task.result() - assert result.is_error - assert "does not support task-augmented execution" in str(result) - - # Prompt with mode="forbidden" raises MCPError when called with task=True - with pytest.raises(MCPError): - await client.get_prompt("my_prompt", task=True) - - # Resource with mode="forbidden" raises MCPError when called with task=True - with pytest.raises(MCPError): - await client.read_resource("test://resource", task=True) + result = await _opted_in_call(mcp, "my_tool") + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "tool result"} -async def test_server_tasks_none_defaults_to_false(): - """Server with tasks=None (or omitted) defaults to False.""" - mcp = FastMCP("test") # tasks=None, defaults to False +async def test_server_tasks_none_defaults_to_forbidden(): + """A server with tasks omitted defaults tools to forbidden (runs sync).""" + mcp = FastMCP("test") # tasks omitted -> forbidden default - @mcp.tool() + @mcp.tool async def my_tool() -> str: return "tool result" - async with Client(mcp, mode="legacy") as client: - # Tool should NOT support background execution (mode="forbidden" from default) - tool_task = await client.call_tool("my_tool", task=True, raise_on_error=False) - assert tool_task.returned_immediately - result = await tool_task.result() - assert result.is_error - assert "does not support task-augmented execution" in str(result) + result = await _opted_in_call(mcp, "my_tool") + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "tool result"} -async def test_component_explicit_false_overrides_server_true(): - """Component with task=False overrides server default of tasks=True.""" - mcp = FastMCP("test", tasks=True) - - @mcp.tool(task=False) - async def no_task_tool() -> str: - return "immediate result" - - @mcp.tool() - async def default_tool() -> str: - return "background result" - - async with Client(mcp, mode="legacy") as client: - # Verify docket registration matches task settings (prefixed keys) - docket = mcp.docket - assert docket is not None - assert ( - "tool:no_task_tool@" not in docket.tasks - ) # task=False means not registered - assert "tool:default_tool@" in docket.tasks # Inherits tasks=True - - # Explicit False (mode="forbidden") returns error when called with task=True - no_task = await client.call_tool( - "no_task_tool", task=True, raise_on_error=False - ) - assert no_task.returned_immediately - result = await no_task.result() - assert result.is_error - assert "does not support task-augmented execution" in str(result) - - # Default should support background execution - default_task = await client.call_tool("default_tool", task=True) - assert not default_task.returned_immediately - - -async def test_component_explicit_true_overrides_server_false(): - """Component with task=True overrides server default of tasks=False.""" +async def test_per_tool_true_overrides_server_false(): + """A per-tool task=True overrides the server default of tasks=False.""" mcp = FastMCP("test", tasks=False) + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def task_tool() -> str: return "background result" - @mcp.tool() + @mcp.tool async def default_tool() -> str: return "immediate result" - async with Client(mcp, mode="legacy") as client: - # Verify docket registration matches task settings (prefixed keys) - docket = mcp.docket - assert docket is not None - assert "tool:task_tool@" in docket.tasks # task=True means registered - assert "tool:default_tool@" not in docket.tasks # Inherits tasks=False + async with running_task_server(mcp): + created = await submit_task(mcp, "task_tool") + assert isinstance(created, CreateTaskResult) - # Explicit True should support background execution despite server default - task = await client.call_tool("task_tool", task=True) - assert not task.returned_immediately - - # Default (mode="forbidden") returns error when called with task=True - default = await client.call_tool( - "default_tool", task=True, raise_on_error=False - ) - assert default.returned_immediately - result = await default.result() - assert result.is_error + # The inherited-forbidden tool still runs synchronously despite the opt-in. + result = await _opted_in_call(mcp, "default_tool") + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "immediate result"} -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_mixed_explicit_and_inherited(): - """Mix of explicit True/False/None on different components.""" - import pytest - from mcp.shared.exceptions import MCPError - - mcp = FastMCP("test", tasks=True) # Server default is True - - @mcp.tool() - async def inherited_tool() -> str: - return "inherits True" - - @mcp.tool(task=True) - async def explicit_true_tool() -> str: - return "explicit True" +async def test_per_tool_false_overrides_server_true(): + """A per-tool task=False overrides the server default of tasks=True.""" + mcp = FastMCP("test", tasks=True) + mcp.add_extension(TasksExtension()) @mcp.tool(task=False) - async def explicit_false_tool() -> str: - return "explicit False" + async def no_task_tool() -> str: + return "immediate result" - @mcp.prompt() - async def inherited_prompt() -> str: - return "inherits True" + @mcp.tool + async def default_tool() -> str: + return "background result" - @mcp.prompt(task=False) - async def explicit_false_prompt() -> str: - return "explicit False" + async with running_task_server(mcp): + # Explicit False runs synchronously even when opted in. + result = await _opted_in_call(mcp, "no_task_tool") + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "immediate result"} - @mcp.resource("test://inherited") - async def inherited_resource() -> str: - return "inherits True" - - @mcp.resource("test://explicit_false", task=False) - async def explicit_false_resource() -> str: - return "explicit False" - - async with Client(mcp, mode="legacy") as client: - # Verify docket registration matches task settings - # Components use prefixed keys: tool:name, prompt:name, resource:uri - docket = mcp.docket - assert docket is not None - # task=True (explicit or inherited) means registered (with prefixed keys) - assert "tool:inherited_tool@" in docket.tasks - assert "tool:explicit_true_tool@" in docket.tasks - assert "prompt:inherited_prompt@" in docket.tasks - assert "resource:test://inherited@" in docket.tasks - # task=False means NOT registered - assert "tool:explicit_false_tool@" not in docket.tasks - assert "prompt:explicit_false_prompt@" not in docket.tasks - assert "resource:test://explicit_false@" not in docket.tasks - - # Tools - inherited = await client.call_tool("inherited_tool", task=True) - assert not inherited.returned_immediately - - explicit_true = await client.call_tool("explicit_true_tool", task=True) - assert not explicit_true.returned_immediately - - # Explicit False (mode="forbidden") returns error - explicit_false = await client.call_tool( - "explicit_false_tool", task=True, raise_on_error=False - ) - assert explicit_false.returned_immediately - result = await explicit_false.result() - assert result.is_error - - # Prompts - inherited_prompt_task = await client.get_prompt("inherited_prompt", task=True) - assert not inherited_prompt_task.returned_immediately - - # Explicit False prompt (mode="forbidden") raises MCPError - with pytest.raises(MCPError): - await client.get_prompt("explicit_false_prompt", task=True) - - # Resources - inherited_resource_task = await client.read_resource( - "test://inherited", task=True - ) - assert not inherited_resource_task.returned_immediately - - # Explicit False resource (mode="forbidden") raises MCPError - with pytest.raises(MCPError): - await client.read_resource("test://explicit_false", task=True) - - -async def test_server_tasks_parameter_sets_component_defaults(): - """Server tasks parameter sets component defaults.""" - # Server tasks=True sets component defaults - mcp = FastMCP("test", tasks=True) - - @mcp.tool() - async def tool_inherits_true() -> str: - return "tool result" - - async with Client(mcp, mode="legacy") as client: - # Tool inherits tasks=True from server - tool_task = await client.call_tool("tool_inherits_true", task=True) - assert not tool_task.returned_immediately - - # Server tasks=False sets component defaults - mcp2 = FastMCP("test2", tasks=False) - - @mcp2.tool() - async def tool_inherits_false() -> str: - return "tool result" - - async with Client(mcp2, mode="legacy") as client: - # Tool inherits tasks=False (mode="forbidden") - returns error - tool_task = await client.call_tool( - "tool_inherits_false", task=True, raise_on_error=False - ) - assert tool_task.returned_immediately - result = await tool_task.result() - assert result.is_error - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_resource_template_inherits_server_tasks_default(): - """Resource templates inherit server tasks default.""" - mcp = FastMCP("test", tasks=True) - - @mcp.resource("test://{item_id}") - async def templated_resource(item_id: str) -> str: - return f"resource {item_id}" - - async with Client(mcp, mode="legacy") as client: - # Template should support background execution - resource_task = await client.read_resource("test://123", task=True) - assert not resource_task.returned_immediately - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_multiple_components_same_name_different_tasks(): - """Different component types with same name can have different task settings.""" - import pytest - from mcp.shared.exceptions import MCPError - - mcp = FastMCP("test", tasks=False) - - @mcp.tool(task=True) - async def shared_name() -> str: - return "tool result" - - @mcp.prompt() - async def shared_name_prompt() -> str: - return "prompt result" - - async with Client(mcp, mode="legacy") as client: - # Tool with explicit True should support background execution - tool_task = await client.call_tool("shared_name", task=True) - assert not tool_task.returned_immediately - - # Prompt inheriting False (mode="forbidden") raises MCPError - with pytest.raises(MCPError): - await client.get_prompt("shared_name_prompt", task=True) + # The inherited-optional tool tasks when opted in. + created = await submit_task(mcp, "default_tool") + assert isinstance(created, CreateTaskResult) async def test_task_with_custom_tool_name(): - """Tools with custom names work correctly as tasks (issue #2642). + """Tools registered under a custom name task correctly (issue #2642). When a tool is registered with a custom name different from the function - name, task execution should use the custom name for Docket lookup. + name, task execution uses the custom name for Docket lookup. """ mcp = FastMCP("test", tasks=True) + mcp.add_extension(TasksExtension()) async def my_function() -> str: return "result from custom-named tool" mcp.tool(my_function, name="custom-tool-name") - async with Client(mcp, mode="legacy") as client: - # Verify the tool is registered with its custom name in Docket (prefixed key) - docket = mcp.docket - assert docket is not None - assert "tool:custom-tool-name@" in docket.tasks - - # Call the tool as a task using its custom name - task = await client.call_tool("custom-tool-name", task=True) - assert not task.returned_immediately - result = await task - assert result.data == "result from custom-named tool" - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_task_with_custom_resource_name(): - """Resources with custom names work correctly as tasks. - - Resources are registered/looked up by their .key (URI), not their name. - """ - mcp = FastMCP("test", tasks=True) - - @mcp.resource("test://resource", name="custom-resource-name") - async def my_resource_func() -> str: - return "result from custom-named resource" - - async with Client(mcp, mode="legacy") as client: - # Verify the resource is registered with its key (prefixed URI) in Docket - docket = mcp.docket - assert docket is not None - assert "resource:test://resource@" in docket.tasks - - # Call the resource as a task - task = await client.read_resource("test://resource", task=True) - assert not task.returned_immediately - result = await task.result() - assert result[0].text == "result from custom-named resource" - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_task_with_custom_template_name(): - """Resource templates with custom names work correctly as tasks. - - Templates are registered/looked up by their .key (uri_template), not their name. - """ - mcp = FastMCP("test", tasks=True) - - @mcp.resource("test://{item_id}", name="custom-template-name") - async def my_template_func(item_id: str) -> str: - return f"result for {item_id}" - - async with Client(mcp, mode="legacy") as client: - # Verify the template is registered with its key (prefixed uri_template) in Docket - docket = mcp.docket - assert docket is not None - assert "template:test://{item_id}@" in docket.tasks - - # Call the template as a task - task = await client.read_resource("test://123", task=True) - assert not task.returned_immediately - result = await task.result() - assert result[0].text == "result for 123" + async with running_task_server(mcp): + final = await run_task(mcp, "custom-tool-name") + assert final.status == "completed" + assert final.result["structuredContent"] == { + "result": "result from custom-named tool" + } diff --git a/tests/tasks/server/test_snapshot_restore.py b/tests/tasks/server/test_snapshot_restore.py index 9e30f0314..532ed7cbb 100644 --- a/tests/tasks/server/test_snapshot_restore.py +++ b/tests/tasks/server/test_snapshot_restore.py @@ -12,7 +12,6 @@ from __future__ import annotations from unittest.mock import patch -import pytest from fastmcp_tasks.context import ( TaskContextSnapshot, _recall_snapshot, @@ -23,45 +22,45 @@ from mcp.server.auth.middleware.auth_context import auth_context_var from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from fastmcp import FastMCP -from fastmcp.client import Client from fastmcp.server.auth import AccessToken from fastmcp.server.dependencies import get_access_token - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + running_task_server, + submit_task, + wait_for_task, ) async def test_snapshot_restored_before_user_code_runs(): """A tool with no declared deps finds the snapshot already cached.""" mcp = FastMCP("snapshot-restore-test") - seen_cached: list[bool] = [] + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) - async def bare_tool() -> str: + async def bare_tool() -> bool: info = get_task_context() assert info is not None - seen_cached.append(_recall_snapshot(info.task_id) is not None) - return "ok" + return _recall_snapshot(info.task_id) is not None - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("bare_tool", {}, task=True) - await task.result() + async with running_task_server(mcp): + created = await submit_task(mcp, "bare_tool", {}) + final = await wait_for_task(mcp, created.task_id) - assert seen_cached == [True] + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": True} async def test_get_access_token_in_bg_task_without_context_dep(): """Issue #3897 repro: get_access_token() works in a bg task that does not declare Context as a dependency.""" mcp = FastMCP("access-token-test") - seen_tokens: list[str | None] = [] + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def bare_tool() -> str: token = get_access_token() - seen_tokens.append(token.token if token else None) - return "ok" + return token.token if token else "no-token" test_token = AccessToken( token="jwt-3897", @@ -71,36 +70,36 @@ async def test_get_access_token_in_bg_task_without_context_dep(): ) auth_context_var.set(AuthenticatedUser(test_token)) - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("bare_tool", {}, task=True) - await task.result() + async with running_task_server(mcp): + created = await submit_task(mcp, "bare_tool", {}) + final = await wait_for_task(mcp, created.task_id) - assert seen_tokens == ["jwt-3897"] + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "jwt-3897"} async def test_restore_failure_is_nonfatal(): """If deserialization blows up, the task still runs to completion and the snapshot cache stays empty.""" mcp = FastMCP("restore-failure-test") - seen_cached: list[bool] = [] + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) - async def bare_tool() -> str: + async def bare_tool() -> bool: info = get_task_context() assert info is not None - seen_cached.append(_recall_snapshot(info.task_id) is not None) - return "ok" + return _recall_snapshot(info.task_id) is not None def boom(*_args, **_kwargs): raise RuntimeError("simulated deserialization failure") - async with Client(mcp, mode="legacy") as client: + async with running_task_server(mcp): with patch.object(TaskContextSnapshot, "from_json", boom): - task = await client.call_tool("bare_tool", {}, task=True) - result = await task.result() + created = await submit_task(mcp, "bare_tool", {}) + final = await wait_for_task(mcp, created.task_id) - assert result.data == "ok" - assert seen_cached == [False] + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": False} async def test_restore_skipped_for_non_fastmcp_task_keys(): diff --git a/tests/tasks/server/test_sync_function_task_disabled.py b/tests/tasks/server/test_sync_function_task_disabled.py index d6147f95f..5230477ec 100644 --- a/tests/tasks/server/test_sync_function_task_disabled.py +++ b/tests/tasks/server/test_sync_function_task_disabled.py @@ -1,21 +1,16 @@ """ Tests that synchronous functions cannot be used as background tasks. -Docket requires async functions for background execution. FastMCP raises -ValueError when task=True is used with a sync function. +SEP-2663 tasks are tools-only. Docket requires async functions for background +execution, so FastMCP raises ValueError when task=True is used with a sync tool +function. These are registration-time checks and need no running server. """ import pytest from fastmcp import FastMCP -from fastmcp.prompts.function_prompt import FunctionPrompt -from fastmcp.resources.function_resource import FunctionResource from fastmcp.tools.function_tool import FunctionTool -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) - async def test_sync_tool_with_explicit_task_true_raises(): """Sync tool with task=True raises ValueError.""" @@ -45,62 +40,6 @@ async def test_sync_tool_with_inherited_task_true_raises(): return x * 2 -async def test_sync_prompt_with_explicit_task_true_raises(): - """Sync prompt with task=True raises ValueError.""" - mcp = FastMCP("test") - - with pytest.raises( - ValueError, match="uses a sync function but has task execution enabled" - ): - - @mcp.prompt(task=True) - def sync_prompt() -> str: - """A synchronous prompt.""" - return "Hello" - - -async def test_sync_prompt_with_inherited_task_true_raises(): - """Sync prompt inheriting task=True from server raises ValueError.""" - mcp = FastMCP("test", tasks=True) - - with pytest.raises( - ValueError, match="uses a sync function but has task execution enabled" - ): - - @mcp.prompt() # Inherits task=True from server - def sync_prompt() -> str: - """A synchronous prompt.""" - return "Hello" - - -async def test_sync_resource_with_explicit_task_true_raises(): - """Sync resource with task=True raises ValueError.""" - mcp = FastMCP("test") - - with pytest.raises( - ValueError, match="uses a sync function but has task execution enabled" - ): - - @mcp.resource("test://sync", task=True) - def sync_resource() -> str: - """A synchronous resource.""" - return "data" - - -async def test_sync_resource_with_inherited_task_true_raises(): - """Sync resource inheriting task=True from server raises ValueError.""" - mcp = FastMCP("test", tasks=True) - - with pytest.raises( - ValueError, match="uses a sync function but has task execution enabled" - ): - - @mcp.resource("test://sync") # Inherits task=True from server - def sync_resource() -> str: - """A synchronous resource.""" - return "data" - - async def test_async_tool_with_task_true_remains_enabled(): """Async tools with task=True keep task support enabled.""" mcp = FastMCP("test") @@ -110,42 +49,11 @@ async def test_async_tool_with_task_true_remains_enabled(): """An async tool.""" return x * 2 - # Tool should have task mode="optional" and be a FunctionTool tool = await mcp.get_tool("async_tool") assert isinstance(tool, FunctionTool) assert tool.task_config.mode == "optional" -async def test_async_prompt_with_task_true_remains_enabled(): - """Async prompts with task=True keep task support enabled.""" - mcp = FastMCP("test") - - @mcp.prompt(task=True) - async def async_prompt() -> str: - """An async prompt.""" - return "Hello" - - # Prompt should have task mode="optional" and be a FunctionPrompt - prompt = await mcp.get_prompt("async_prompt") - assert isinstance(prompt, FunctionPrompt) - assert prompt.task_config.mode == "optional" - - -async def test_async_resource_with_task_true_remains_enabled(): - """Async resources with task=True keep task support enabled.""" - mcp = FastMCP("test") - - @mcp.resource("test://async", task=True) - async def async_resource() -> str: - """An async resource.""" - return "data" - - # Resource should have task mode="optional" and be a FunctionResource - resource = await mcp.get_resource("test://async") - assert isinstance(resource, FunctionResource) - assert resource.task_config.mode == "optional" - - async def test_sync_tool_with_task_false_works(): """Sync tool with explicit task=False works (no error).""" mcp = FastMCP("test", tasks=True) @@ -160,36 +68,8 @@ async def test_sync_tool_with_task_false_works(): assert tool.task_config.mode == "forbidden" -async def test_sync_prompt_with_task_false_works(): - """Sync prompt with explicit task=False works (no error).""" - mcp = FastMCP("test", tasks=True) - - @mcp.prompt(task=False) # Explicitly disable - def sync_prompt() -> str: - """A synchronous prompt.""" - return "Hello" - - prompt = await mcp.get_prompt("sync_prompt") - assert isinstance(prompt, FunctionPrompt) - assert prompt.task_config.mode == "forbidden" - - -async def test_sync_resource_with_task_false_works(): - """Sync resource with explicit task=False works (no error).""" - mcp = FastMCP("test", tasks=True) - - @mcp.resource("test://sync", task=False) # Explicitly disable - def sync_resource() -> str: - """A synchronous resource.""" - return "data" - - resource = await mcp.get_resource("test://sync") - assert isinstance(resource, FunctionResource) - assert resource.task_config.mode == "forbidden" - - # ============================================================================= -# Callable classes and staticmethods with async __call__ +# Callable classes with async __call__ # ============================================================================= @@ -201,24 +81,10 @@ async def test_async_callable_class_tool_with_task_true_works(): async def __call__(self, x: int) -> int: return x * 2 - # Callable classes use Tool.from_function() directly tool = Tool.from_function(AsyncCallableTool(), task=True) assert tool.task_config.mode == "optional" -async def test_async_callable_class_prompt_with_task_true_works(): - """Callable class with async __call__ and task=True should work.""" - from fastmcp.prompts import Prompt - - class AsyncCallablePrompt: - async def __call__(self) -> str: - return "Hello" - - # Callable classes use Prompt.from_function() directly - prompt = Prompt.from_function(AsyncCallablePrompt(), task=True) - assert prompt.task_config.mode == "optional" - - async def test_sync_callable_class_tool_with_task_true_raises(): """Callable class with sync __call__ and task=True should raise.""" from fastmcp.tools import Tool diff --git a/tests/tasks/server/test_task_capabilities.py b/tests/tasks/server/test_task_capabilities.py index a79bd753d..8467ccc42 100644 --- a/tests/tasks/server/test_task_capabilities.py +++ b/tests/tasks/server/test_task_capabilities.py @@ -1,97 +1,40 @@ -""" -Tests for SEP-1686 task capabilities declaration. +"""Advertisement of the SEP-2663 tasks extension capability. -Verifies that the server correctly advertises task support. -Task protocol is now always enabled. +A server with the tasks extension registered advertises the +`io.modelcontextprotocol/tasks` extension in its capabilities; a server without +it does not. """ -import pytest -from fastmcp_tasks._legacy_wire.capabilities import get_task_capabilities +from __future__ import annotations from fastmcp import FastMCP from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) +from fastmcp.utilities.tasks import TASKS_EXTENSION_ID +from fastmcp_tasks import TasksExtension -async def test_capabilities_include_tasks(): - """Server capabilities always include tasks in first-class field (SEP-1686).""" +async def test_extension_capability_advertised(): + """The tasks extension is advertised when registered.""" + mcp = FastMCP("capability-test") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def my_tool() -> str: + return "ok" + + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert extensions.get(TASKS_EXTENSION_ID) == {} + + +async def test_extension_capability_absent_without_extension(): + """The tasks extension is not advertised when no extension is registered.""" mcp = FastMCP("capability-test") - @mcp.tool() - async def test_tool() -> str: - return "test" + @mcp.tool + async def my_tool() -> str: + return "ok" - async with Client(mcp, mode="legacy") as client: - # Get server initialization result which includes capabilities - init_result = client.initialize_result - - # Verify tasks capability is present as a first-class field (not experimental) - assert init_result.capabilities.tasks is not None - assert init_result.capabilities.tasks == get_task_capabilities() - # Verify it's NOT in experimental - assert "tasks" not in (init_result.capabilities.experimental or {}) - - -def test_only_tools_advertise_task_support(): - """Task requests advertise tools only, not prompts/resources (sdk-feedback #3). - - SDK v2 b1 ``ReadResourceRequestParams`` / ``GetPromptRequestParams`` have no - ``task`` field, so resource/prompt task submissions always graceful-degrade - to synchronous execution. Advertising those capabilities would mislead - clients into sending task-augmented reads/gets, so the honest contract is - tools-only. - """ - capabilities = get_task_capabilities() - assert capabilities is not None - requests = capabilities.requests - assert requests is not None - assert requests.tools is not None - assert requests.tools.call is not None - # No prompt/resource task capability of any form is advertised. - assert getattr(requests, "prompts", None) is None - assert getattr(requests, "resources", None) is None - dumped = requests.model_dump(exclude_none=True) - assert set(dumped) == {"tools"} - - -async def test_client_uses_task_capable_session(): - """Client uses task-capable initialization.""" - mcp = FastMCP("client-cap-test") - - @mcp.tool() - async def test_tool() -> str: - return "test" - - async with Client(mcp, mode="legacy") as client: - # Client should have connected successfully with task capabilities - assert client.initialize_result is not None - # Session should be a ClientSession (task-capable init uses standard session) - assert type(client.session).__name__ == "ClientSession" - - -def test_capabilities_hidden_when_pydocket_too_old(monkeypatch): - """Capability advertisement and handler registration must agree. - - If ``is_docket_available()`` returns False (e.g. an old transitive - pydocket), the server skips registering task handlers — so it must - also stop advertising task capabilities, or clients would discover - task support and then hit "method not found" at runtime. - """ - import importlib.metadata - - from fastmcp.server import dependencies - - original_version = importlib.metadata.version - - def fake_version(name: str) -> str: - if name == "pydocket": - return "0.16.6" - return original_version(name) - - monkeypatch.setattr(dependencies, "_DOCKET_AVAILABLE", None) - monkeypatch.setattr(importlib.metadata, "version", fake_version) - - assert get_task_capabilities() is None + async with Client(mcp, mode="auto") as client: + extensions = client.server_capabilities.extensions or {} + assert TASKS_EXTENSION_ID not in extensions diff --git a/tests/tasks/server/test_task_config.py b/tests/tasks/server/test_task_config.py index ba9cc8cb7..2f687da77 100644 --- a/tests/tasks/server/test_task_config.py +++ b/tests/tasks/server/test_task_config.py @@ -1,22 +1,40 @@ -"""Tests for TaskConfig (SEP-1686). +"""Tests for TaskConfig (SEP-2663, tools only). Tests for TaskConfig: -- Mode enforcement (forbidden, optional, required) +- Normalization of boolean task values to TaskConfig +- Sync-function validation +- Tool mode enforcement (forbidden, optional, required) +- Tool execution metadata (task_support in tools/list) - Poll interval configuration """ from datetime import timedelta import pytest +from fastmcp_tasks.models import ( + MISSING_REQUIRED_CLIENT_CAPABILITY, + CreateTaskResult, +) from mcp.shared.exceptions import MCPError -from mcp_types import TextContent, ToolExecution -from mcp_types import Tool as MCPTool +from mcp_types import ToolExecution from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.exceptions import ToolError from fastmcp.tools.base import Tool from fastmcp.utilities.tasks import TaskConfig +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + _opted_in_request, + auth_scope, + call_tool_without_optin, + running_task_server, + submit_task, +) + + +async def _opted_in_call(server: FastMCP, name: str, arguments: dict | None = None): + """Run a `tools/call` WITH the tasks opt-in bound (used to prove sync paths).""" + with auth_scope(None), _opted_in_request(name, arguments or {}, None): + return await server.call_tool(name, arguments or {}) class TestTaskConfigNormalization: @@ -83,248 +101,114 @@ class TestTaskConfigNormalization: assert tool2.task_config.mode == "optional" -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") class TestToolModeEnforcement: - """Test mode enforcement for tools.""" + """Test mode enforcement for tools under the SEP-2663 interceptor.""" - @pytest.fixture - def server(self): - """Create server with tools in different modes.""" + def _server(self) -> FastMCP: mcp = FastMCP("test", tasks=False) + mcp.add_extension(TasksExtension()) @mcp.tool(task=TaskConfig(mode="required")) async def required_tool() -> str: - """Tool that requires task execution.""" return "required result" @mcp.tool(task=TaskConfig(mode="forbidden")) async def forbidden_tool() -> str: - """Tool that forbids task execution.""" return "forbidden result" @mcp.tool(task=TaskConfig(mode="optional")) async def optional_tool() -> str: - """Tool that supports both modes.""" return "optional result" return mcp - async def test_required_mode_without_task_returns_error(self, server): - """Required mode raises error when called without task metadata.""" - async with Client(server, mode="legacy") as client: - with pytest.raises(ToolError) as exc_info: - await client.call_tool("required_tool", {}) - - assert "requires task-augmented execution" in str(exc_info.value) - - async def test_required_mode_with_task_succeeds(self, server): - """Required mode succeeds when called with task metadata.""" - async with Client(server, mode="legacy") as client: - task = await client.call_tool("required_tool", {}, task=True) - assert task is not None - result = await task.result() - assert result.data == "required result" - - async def test_forbidden_mode_with_task_returns_error(self, server): - """Forbidden mode returns error when called with task metadata.""" - async with Client(server, mode="legacy") as client: - # Call with task=True should fail - task = await client.call_tool( - "forbidden_tool", {}, task=True, raise_on_error=False - ) - assert task is not None - # The task should have returned immediately with an error - assert task.returned_immediately - result = await task.result() - # Check for error in the result - assert result.is_error - - async def test_forbidden_mode_without_task_succeeds(self, server): - """Forbidden mode succeeds when called without task metadata.""" - async with Client(server, mode="legacy") as client: - result = await client.call_tool("forbidden_tool", {}) - assert "forbidden result" in str(result) - - async def test_optional_mode_without_task_succeeds(self, server): - """Optional mode succeeds when called without task metadata.""" - async with Client(server, mode="legacy") as client: - result = await client.call_tool("optional_tool", {}) - assert "optional result" in str(result) - - async def test_optional_mode_with_task_succeeds(self, server): - """Optional mode succeeds when called with task metadata.""" - async with Client(server, mode="legacy") as client: - task = await client.call_tool("optional_tool", {}, task=True) - assert task is not None - result = await task.result() - assert result.data == "optional result" - - -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") -class TestResourceModeEnforcement: - """Test mode enforcement for resources.""" - - @pytest.fixture - def server(self): - """Create server with resources in different modes.""" - mcp = FastMCP("test", tasks=False) - - @mcp.resource("resource://required", task=TaskConfig(mode="required")) - async def required_resource() -> str: - """Resource that requires task execution.""" - return "required content" - - @mcp.resource("resource://forbidden", task=TaskConfig(mode="forbidden")) - async def forbidden_resource() -> str: - """Resource that forbids task execution.""" - return "forbidden content" - - @mcp.resource("resource://optional", task=TaskConfig(mode="optional")) - async def optional_resource() -> str: - """Resource that supports both modes.""" - return "optional content" - - return mcp - - async def test_required_resource_without_task_returns_error(self, server): - """Required mode returns error when read without task metadata.""" - from mcp_types import METHOD_NOT_FOUND - - async with Client(server, mode="legacy") as client: + async def test_required_mode_without_opt_in_raises(self): + """Required mode raises -32003 when called without a tasks opt-in.""" + mcp = self._server() + async with running_task_server(mcp): with pytest.raises(MCPError) as exc_info: - await client.read_resource("resource://required") + await call_tool_without_optin(mcp, "required_tool") + assert exc_info.value.error.code == MISSING_REQUIRED_CLIENT_CAPABILITY - assert exc_info.value.error.code == METHOD_NOT_FOUND - assert "requires task-augmented execution" in exc_info.value.error.message + async def test_required_mode_with_opt_in_tasks(self): + """Required mode tasks when the caller opts in.""" + mcp = self._server() + async with running_task_server(mcp): + created = await submit_task(mcp, "required_tool") + assert isinstance(created, CreateTaskResult) - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_required_resource_with_task_succeeds(self, server): - """Required mode succeeds when read with task metadata.""" - async with Client(server, mode="legacy") as client: - task = await client.read_resource("resource://required", task=True) - assert task is not None - result = await task.result() - # Result is a list of resource contents - assert "required content" in str(result) + async def test_forbidden_mode_never_tasks_even_with_opt_in(self): + """Forbidden mode runs synchronously even when the caller opts in.""" + mcp = self._server() + async with running_task_server(mcp): + result = await _opted_in_call(mcp, "forbidden_tool") + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "forbidden result"} - async def test_forbidden_resource_without_task_succeeds(self, server): - """Forbidden mode succeeds when read without task metadata.""" - async with Client(server, mode="legacy") as client: - result = await client.read_resource("resource://forbidden") - assert "forbidden content" in str(result) + async def test_optional_mode_without_opt_in_runs_sync(self): + """Optional mode runs synchronously without a tasks opt-in.""" + mcp = self._server() + async with running_task_server(mcp): + result = await call_tool_without_optin(mcp, "optional_tool") + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "optional result"} + + async def test_optional_mode_with_opt_in_tasks(self): + """Optional mode tasks when the caller opts in.""" + mcp = self._server() + async with running_task_server(mcp): + created = await submit_task(mcp, "optional_tool") + assert isinstance(created, CreateTaskResult) -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") -class TestPromptModeEnforcement: - """Test mode enforcement for prompts.""" - - @pytest.fixture - def server(self): - """Create server with prompts in different modes.""" - mcp = FastMCP("test", tasks=False) - - @mcp.prompt(task=TaskConfig(mode="required")) - async def required_prompt() -> str: - """Prompt that requires task execution.""" - return "required message" - - @mcp.prompt(task=TaskConfig(mode="forbidden")) - async def forbidden_prompt() -> str: - """Prompt that forbids task execution.""" - return "forbidden message" - - @mcp.prompt(task=TaskConfig(mode="optional")) - async def optional_prompt() -> str: - """Prompt that supports both modes.""" - return "optional message" - - return mcp - - async def test_required_prompt_without_task_returns_error(self, server): - """Required mode returns error when called without task metadata.""" - from mcp_types import METHOD_NOT_FOUND - - async with Client(server, mode="legacy") as client: - with pytest.raises(MCPError) as exc_info: - await client.get_prompt("required_prompt") - - assert exc_info.value.error.code == METHOD_NOT_FOUND - assert "requires task-augmented execution" in exc_info.value.error.message - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_required_prompt_with_task_succeeds(self, server): - """Required mode succeeds when called with task metadata.""" - async with Client(server, mode="legacy") as client: - task = await client.get_prompt("required_prompt", task=True) - assert task is not None - result = await task.result() - # Result contains the prompt messages - assert "required message" in str(result) - - async def test_forbidden_prompt_without_task_succeeds(self, server): - """Forbidden mode succeeds when called without task metadata.""" - async with Client(server, mode="legacy") as client: - result = await client.get_prompt("forbidden_prompt") - assert isinstance(result.messages[0].content, TextContent) - assert "forbidden message" in str(result.messages[0].content) - - -@pytest.mark.skip(reason="Phase 3: requires TasksExtension (SEP-2663 adapter)") class TestToolExecutionMetadata: - """Test that ToolExecution.task_support is set correctly in tool metadata.""" + """Test that ToolExecution.task_support is set correctly in tool metadata. + + The tools/list payload is produced by ``Tool.to_mcp_tool()``; these tests + assert on that serialization directly, which is what a server advertises on + the wire. (The FastMCP client session does not yet surface ``execution`` back + to callers, so a client round-trip cannot observe it until Phase 4.) + """ async def test_optional_tool_exposes_task_support(self): - """Tools with task enabled should expose taskSupport in metadata.""" + """Tools with mode=optional expose task_support='optional'.""" mcp = FastMCP("test", tasks=False) + mcp.add_extension(TasksExtension()) @mcp.tool(task=TaskConfig(mode="optional")) async def my_tool() -> str: return "ok" - async with Client(mcp, mode="legacy") as client: - tools = await client.list_tools() - tool = next(t for t in tools if t.name == "my_tool") - assert isinstance(tool, MCPTool) - assert isinstance(tool.execution, ToolExecution) - assert tool.execution.task_support == "optional" + tool = await mcp.get_tool("my_tool") + execution = tool.to_mcp_tool().execution + assert isinstance(execution, ToolExecution) + assert execution.task_support == "optional" async def test_required_tool_exposes_task_support(self): - """Tools with mode=required should expose task_support='required'.""" + """Tools with mode=required expose task_support='required'.""" mcp = FastMCP("test", tasks=False) + mcp.add_extension(TasksExtension()) @mcp.tool(task=TaskConfig(mode="required")) async def my_tool() -> str: return "ok" - async with Client(mcp, mode="legacy") as client: - tools = await client.list_tools() - tool = next(t for t in tools if t.name == "my_tool") - assert isinstance(tool, MCPTool) - assert isinstance(tool.execution, ToolExecution) - assert tool.execution.task_support == "required" + tool = await mcp.get_tool("my_tool") + execution = tool.to_mcp_tool().execution + assert isinstance(execution, ToolExecution) + assert execution.task_support == "required" async def test_forbidden_tool_has_no_execution(self): - """Tools with mode=forbidden should not expose execution metadata.""" + """Tools with mode=forbidden do not expose execution metadata.""" mcp = FastMCP("test", tasks=False) + mcp.add_extension(TasksExtension()) @mcp.tool(task=TaskConfig(mode="forbidden")) async def my_tool() -> str: return "ok" - async with Client(mcp, mode="legacy") as client: - tools = await client.list_tools() - tool = next(t for t in tools if t.name == "my_tool") - assert tool.execution is None + tool = await mcp.get_tool("my_tool") + assert tool.to_mcp_tool().execution is None class TestSyncFunctionValidation: diff --git a/tests/tasks/server/test_task_dependencies.py b/tests/tasks/server/test_task_dependencies.py index 87304963d..e40770d46 100644 --- a/tests/tasks/server/test_task_dependencies.py +++ b/tests/tasks/server/test_task_dependencies.py @@ -1,10 +1,15 @@ """Tests for dependency injection in background tasks. -These tests verify that Docket's dependency system works correctly when -user functions are queued as background tasks. Dependencies like CurrentDocket(), +These tests verify that Docket's dependency system works correctly when tool +functions are queued as background tasks. Dependencies like CurrentDocket(), CurrentFastMCP(), and Depends() should be resolved in the worker context. + +SEP-2663 is tools-only, so only tools carry a task-capable config; the removed +prompt/resource task cases are gone. """ +from __future__ import annotations + from contextlib import asynccontextmanager from typing import Any, cast @@ -13,32 +18,30 @@ from fastmcp_tasks.dependencies import CurrentDocket from uncalled_for import Depends from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.exceptions import ToolError from fastmcp.server.dependencies import CurrentFastMCP - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + call_tool_without_optin, + run_task, + running_task_server, ) @pytest.fixture -async def dependency_server(): - """Create a FastMCP server with dependency-using background tasks.""" +def dependency_server() -> FastMCP: + """A FastMCP server with dependency-using background tools.""" mcp = FastMCP("dependency-test-server") + mcp.add_extension(TasksExtension()) - # Track dependency injection - injected_values = [] + injected_values: list[tuple[str, Any]] = [] @mcp.tool(task=True) async def tool_with_docket_dependency(docket=CurrentDocket()) -> str: - """Background tool that uses CurrentDocket dependency.""" injected_values.append(("docket", docket)) return f"Docket: {docket is not None}" @mcp.tool(task=True) async def tool_with_server_dependency(server=CurrentFastMCP()) -> str: - """Background tool that uses CurrentFastMCP dependency.""" injected_values.append(("server", server)) return f"Server: {server.name}" @@ -46,7 +49,6 @@ async def dependency_server(): async def tool_with_custom_dependency( value: int, multiplier: int = Depends(lambda: 10) ) -> int: - """Background tool with custom Depends().""" injected_values.append(("multiplier", multiplier)) return value * multiplier @@ -56,188 +58,108 @@ async def dependency_server(): docket=CurrentDocket(), server=CurrentFastMCP(), ) -> str: - """Background tool with multiple dependencies.""" injected_values.append(("multi_docket", docket)) injected_values.append(("multi_server", server)) return f"{name} on {server.name}" - @mcp.prompt(task=True) - async def prompt_with_server_dependency(topic: str, server=CurrentFastMCP()) -> str: - """Background prompt that uses CurrentFastMCP dependency.""" - injected_values.append(("prompt_server", server)) - return f"Prompt from {server.name} about {topic}" - - @mcp.resource("file://data.txt", task=True) - async def resource_with_docket_dependency(docket=CurrentDocket()) -> str: - """Background resource that uses CurrentDocket dependency.""" - injected_values.append(("resource_docket", docket)) - return f"Resource via Docket: {docket is not None}" - - # Expose for test assertions mcp._injected_values = injected_values # type: ignore[attr-defined] # ty:ignore[unresolved-attribute] return mcp async def test_background_tool_receives_docket_dependency(dependency_server): - """Background tools can use CurrentDocket() and it resolves correctly.""" - async with Client(dependency_server, mode="legacy") as client: - task = await client.call_tool("tool_with_docket_dependency", {}, task=True) + """Background tools can use CurrentDocket() and it resolves in the worker.""" + async with running_task_server(dependency_server): + final = await run_task(dependency_server, "tool_with_docket_dependency", {}) - # Verify it's background - assert not task.returned_immediately - - # Get result - will execute in Docket worker - result = await task - - # Verify dependency was injected - assert len(dependency_server._injected_values) == 1 - dep_type, dep_value = dependency_server._injected_values[0] - assert dep_type == "docket" - assert dep_value is not None - assert "Docket: True" in result.data + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "Docket: True"} + assert len(dependency_server._injected_values) == 1 + dep_type, dep_value = dependency_server._injected_values[0] + assert dep_type == "docket" + assert dep_value is not None async def test_background_tool_receives_server_dependency(dependency_server): - """Background tools can use CurrentFastMCP() and get the actual FastMCP server.""" + """Background tools can use CurrentFastMCP() and get the actual server.""" dependency_server._injected_values.clear() - async with Client(dependency_server, mode="legacy") as client: - task = await client.call_tool("tool_with_server_dependency", {}, task=True) + async with running_task_server(dependency_server): + final = await run_task(dependency_server, "tool_with_server_dependency", {}) - # Verify background execution - assert not task.returned_immediately - - result = await task - - # Check the server instance was injected - assert len(dependency_server._injected_values) == 1 - dep_type, dep_value = dependency_server._injected_values[0] - assert dep_type == "server" - assert dep_value is dependency_server # Same instance! - assert f"Server: {dependency_server.name}" in result.data + assert final.status == "completed" + assert final.result["structuredContent"] == { + "result": f"Server: {dependency_server.name}" + } + assert len(dependency_server._injected_values) == 1 + dep_type, dep_value = dependency_server._injected_values[0] + assert dep_type == "server" + assert dep_value is dependency_server # Same instance! async def test_background_tool_receives_custom_depends(dependency_server): """Background tools can use Depends() with custom functions.""" dependency_server._injected_values.clear() - async with Client(dependency_server, mode="legacy") as client: - task = await client.call_tool( - "tool_with_custom_dependency", {"value": 5}, task=True + async with running_task_server(dependency_server): + final = await run_task( + dependency_server, "tool_with_custom_dependency", {"value": 5} ) - assert not task.returned_immediately - - result = await task - - # Check dependency was resolved - assert len(dependency_server._injected_values) == 1 - dep_type, dep_value = dependency_server._injected_values[0] - assert dep_type == "multiplier" - assert dep_value == 10 - assert result.data == 50 # 5 * 10 + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": 50} # 5 * 10 + assert len(dependency_server._injected_values) == 1 + dep_type, dep_value = dependency_server._injected_values[0] + assert dep_type == "multiplier" + assert dep_value == 10 async def test_background_tool_with_multiple_dependencies(dependency_server): - """Background tools can have multiple dependencies injected simultaneously.""" + """Background tools can have multiple dependencies injected at once.""" dependency_server._injected_values.clear() - async with Client(dependency_server, mode="legacy") as client: - task = await client.call_tool( - "tool_with_multiple_dependencies", {"name": "test"}, task=True + async with running_task_server(dependency_server): + final = await run_task( + dependency_server, "tool_with_multiple_dependencies", {"name": "test"} ) - assert not task.returned_immediately + assert final.status == "completed" + assert final.result["structuredContent"] == { + "result": f"test on {dependency_server.name}" + } - await task + dep_types = {item[0] for item in dependency_server._injected_values} + assert "multi_docket" in dep_types + assert "multi_server" in dep_types - # Both dependencies should be injected - assert len(dependency_server._injected_values) == 2 - - dep_types = {item[0] for item in dependency_server._injected_values} - assert "multi_docket" in dep_types - assert "multi_server" in dep_types - - # Verify values - server_dep = next( - v for t, v in dependency_server._injected_values if t == "multi_server" - ) - assert server_dep is dependency_server - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_background_prompt_receives_dependencies(dependency_server): - """Background prompts can use dependency injection.""" - dependency_server._injected_values.clear() - - async with Client(dependency_server, mode="legacy") as client: - task = await client.get_prompt( - "prompt_with_server_dependency", {"topic": "AI"}, task=True - ) - - assert not task.returned_immediately - - await task - - # Check dependency was injected - assert len(dependency_server._injected_values) == 1 - dep_type, dep_value = dependency_server._injected_values[0] - assert dep_type == "prompt_server" - assert dep_value is dependency_server - - -@pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, -) -async def test_background_resource_receives_dependencies(dependency_server): - """Background resources can use dependency injection.""" - dependency_server._injected_values.clear() - - async with Client(dependency_server, mode="legacy") as client: - task = await client.read_resource("file://data.txt", task=True) - - assert not task.returned_immediately - - await task - - # Check dependency was injected - assert len(dependency_server._injected_values) == 1 - dep_type, dep_value = dependency_server._injected_values[0] - assert dep_type == "resource_docket" - assert dep_value is not None + server_dep = next( + v for t, v in dependency_server._injected_values if t == "multi_server" + ) + assert server_dep is dependency_server async def test_foreground_tool_dependencies_unaffected(dependency_server): - """Synchronous tools (task=False) still get dependencies as before.""" + """Synchronous tools still get their dependencies as before.""" dependency_server._injected_values.clear() - @dependency_server.tool() # task=False + @dependency_server.tool async def sync_tool(server=CurrentFastMCP()) -> str: dependency_server._injected_values.append(("sync_server", server)) return f"Sync: {server.name}" - async with Client(dependency_server, mode="legacy") as client: - await client.call_tool("sync_tool", {}) + async with running_task_server(dependency_server): + await call_tool_without_optin(dependency_server, "sync_tool", {}) - # Should execute immediately - assert len(dependency_server._injected_values) == 1 - assert dependency_server._injected_values[0][1] is dependency_server + assert len(dependency_server._injected_values) == 1 + assert dependency_server._injected_values[0][1] is dependency_server async def test_dependency_context_managers_cleaned_up_in_background(): - """Context manager dependencies are properly cleaned up after background task.""" - cleanup_called = [] + """Context-manager dependencies are cleaned up after a background task.""" + cleanup_called: list[str] = [] mcp = FastMCP("cleanup-test") + mcp.add_extension(TasksExtension()) @asynccontextmanager async def tracked_connection(): @@ -254,18 +176,18 @@ async def test_dependency_context_managers_cleaned_up_in_background(): assert "exit" not in cleanup_called # Still open during execution return f"Used: {conn}" - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("use_connection", {"name": "test"}, task=True) - result = await task + async with running_task_server(mcp): + final = await run_task(mcp, "use_connection", {"name": "test"}) - # After task completes, cleanup should have been called - assert cleanup_called == ["enter", "exit"] - assert "Used: connection" in result.data + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "Used: connection"} + assert cleanup_called == ["enter", "exit"] async def test_dependency_errors_propagate_to_task_failure(): """If dependency resolution fails, the background task should fail.""" mcp = FastMCP("error-test") + mcp.add_extension(TasksExtension()) async def failing_dependency(): raise ValueError("Dependency failed!") @@ -276,15 +198,8 @@ async def test_dependency_errors_propagate_to_task_failure(): ) -> str: return f"Got: {dep}" - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool( - "tool_with_failing_dep", {"value": "test"}, task=True - ) + async with running_task_server(mcp): + final = await run_task(mcp, "tool_with_failing_dep", {"value": "test"}) - # Task should fail due to dependency error - with pytest.raises(ToolError, match="Failed to resolve dependencies"): - await task.result() - - # Verify it reached failed state - status = await task.status() - assert status.status == "failed" + assert final.status == "failed" + assert final.error is not None diff --git a/tests/tasks/server/test_task_elicitation_relay.py b/tests/tasks/server/test_task_elicitation_relay.py index 770b28801..a36e202fe 100644 --- a/tests/tasks/server/test_task_elicitation_relay.py +++ b/tests/tasks/server/test_task_elicitation_relay.py @@ -1,196 +1,217 @@ -"""Tests for background task elicitation relay (notifications.py). +"""In-task elicitation under SEP-2663 (poll-based input). -The relay bridges distributed background tasks to clients via the standard -MCP elicitation/create protocol. When a worker calls ctx.elicit(), the -notification subscriber detects the input_required notification and sends -an elicitation/create request to the client session. The client's -elicitation_handler fires, and the relay pushes the response to Redis -for the blocked worker. - -These tests use Client(mcp, mode="legacy") with the real memory:// Docket backend. +A background worker that calls ``ctx.elicit()`` has no live request, so SEP-2663 +parks the request and the task's ``tasks/get`` status flips to ``input_required`` +with the outstanding ``inputRequests``. The caller answers with ``tasks/update`` +and the parked worker resumes. This replaces the SEP-1686 push relay (which sent +``elicitation/create`` over a back-channel); the accept/decline/cancel semantics, +structured round-trips, and sequential elicitations are preserved, driven here +in-process because there is no client task API until Phase 4. """ +from __future__ import annotations + import asyncio from dataclasses import dataclass +from typing import Any -import pytest +import fastmcp_tasks.input_store as input_store from pydantic import BaseModel from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.client.elicitation import ElicitResult from fastmcp.server.context import Context from fastmcp.server.elicitation import ( AcceptedElicitation, CancelledElicitation, DeclinedElicitation, ) - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + get_task, + running_task_server, + submit_task, + update_task, + wait_for_task, ) -class TestElicitationRelay: - """E2E tests for elicitation flowing through the standard MCP protocol.""" - - async def test_accept_via_elicitation_handler(self): - """Tool elicits, client handler accepts, tool gets the value.""" - mcp = FastMCP("relay-accept") - - @mcp.tool(task=True) - async def ask_name(ctx: Context) -> str: - result = await ctx.elicit("What is your name?", str) - if isinstance(result, AcceptedElicitation): - return f"Hello, {result.data}!" - return "No name" - - async def handler(message, response_type, params, ctx): - assert message == "What is your name?" - return ElicitResult(action="accept", content={"value": "Alice"}) - - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("ask_name", {}, task=True) - result = await task.result() - assert result.data == "Hello, Alice!" - - async def test_decline_via_elicitation_handler(self): - """Tool elicits, client handler declines, tool gets DeclinedElicitation.""" - mcp = FastMCP("relay-decline") - - @mcp.tool(task=True) - async def optional_input(ctx: Context) -> str: - result = await ctx.elicit("Provide a name?", str) - if isinstance(result, DeclinedElicitation): - return "User declined" - if isinstance(result, AcceptedElicitation): - return f"Got: {result.data}" - return "Cancelled" - - async def handler(message, response_type, params, ctx): - return ElicitResult(action="decline") - - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("optional_input", {}, task=True) - result = await task.result() - assert result.data == "User declined" - - async def test_cancel_via_elicitation_handler(self): - """Tool elicits, client handler cancels, tool gets CancelledElicitation.""" - mcp = FastMCP("relay-cancel") - - @mcp.tool(task=True) - async def cancellable(ctx: Context) -> str: - result = await ctx.elicit("Input?", str) - if isinstance(result, CancelledElicitation): - return "Cancelled" - return "Not cancelled" - - async def handler(message, response_type, params, ctx): - return ElicitResult(action="cancel") - - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("cancellable", {}, task=True) - result = await task.result() - assert result.data == "Cancelled" - - async def test_dataclass_round_trips_through_relay(self): - """Structured dataclass type round-trips through the relay.""" - mcp = FastMCP("relay-dataclass") - - @dataclass - class UserInfo: - name: str - age: int - - @mcp.tool(task=True) - async def get_user(ctx: Context) -> str: - 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 def handler(message, response_type, params, ctx): - return ElicitResult(action="accept", content={"name": "Bob", "age": 30}) - - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("get_user", {}, task=True) - result = await task.result() - assert result.data == "Bob is 30" - - async def test_pydantic_model_round_trips_through_relay(self): - """Structured Pydantic model round-trips through the relay.""" - mcp = FastMCP("relay-pydantic") - - class Config(BaseModel): - host: str - port: int - - @mcp.tool(task=True) - async def get_config(ctx: Context) -> str: - result = await ctx.elicit("Server config?", Config) - if isinstance(result, AcceptedElicitation): - assert isinstance(result.data, Config) - return f"{result.data.host}:{result.data.port}" - return "No config" - - async def handler(message, response_type, params, ctx): - return ElicitResult( - action="accept", content={"host": "localhost", "port": 8080} +async def _wait_for_input_required(server: FastMCP, task_id: str, timeout: float = 5.0): + """Poll until the task is waiting on input, returning the GetTaskResult.""" + deadline = asyncio.get_event_loop().time() + timeout + while True: + got = await get_task(server, task_id) + if got.status == "input_required": + return got + if got.status in ("completed", "failed", "cancelled"): + raise AssertionError( + f"Task {task_id} reached {got.status!r} before requesting input" ) + if asyncio.get_event_loop().time() >= deadline: + raise TimeoutError(f"Task {task_id} never requested input") + await asyncio.sleep(0.02) - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("get_config", {}, task=True) - result = await task.result() - assert result.data == "localhost:8080" - async def test_multiple_sequential_elicitations(self): - """Tool calls ctx.elicit() twice, both go through the relay.""" - mcp = FastMCP("relay-multi") +async def _drive(server: FastMCP, name: str, answers: list[dict[str, Any]]) -> str: + """Submit a task, answer each elicitation in turn, return its result text.""" + created = await submit_task(server, name, {}) + for answer in answers: + got = await _wait_for_input_required(server, created.task_id) + key = next(iter(got.input_requests)) + request = got.input_requests[key] + assert request["method"] == "elicitation/create" + await update_task(server, created.task_id, {key: answer}) + final = await wait_for_task(server, created.task_id) + assert final.status == "completed", final.error + return final.result["content"][0]["text"] - @mcp.tool(task=True) - async def two_questions(ctx: Context) -> str: - r1 = await ctx.elicit("First name?", str) - r2 = await ctx.elicit("Last name?", str) - if isinstance(r1, AcceptedElicitation) and isinstance( - r2, AcceptedElicitation - ): - return f"{r1.data} {r2.data}" - return "Incomplete" - call_count = 0 +async def test_accept_answers_the_elicitation(): + mcp = FastMCP("relay-accept") + mcp.add_extension(TasksExtension()) - async def handler(message, response_type, params, ctx): - nonlocal call_count - call_count += 1 - if call_count == 1: - assert message == "First name?" - return ElicitResult(action="accept", content={"value": "Jane"}) - else: - assert message == "Last name?" - return ElicitResult(action="accept", content={"value": "Doe"}) + @mcp.tool(task=True) + async def ask_name(ctx: Context) -> str: + result = await ctx.elicit("What is your name?", str) + if isinstance(result, AcceptedElicitation): + return f"Hello, {result.data}!" + return "No name" - async with Client(mcp, mode="legacy", elicitation_handler=handler) as client: - task = await client.call_tool("two_questions", {}, task=True) - result = await task.result() - assert result.data == "Jane Doe" - assert call_count == 2 + async with running_task_server(mcp): + text = await _drive( + mcp, "ask_name", [{"action": "accept", "content": {"value": "Alice"}}] + ) + assert text == "Hello, Alice!" - async def test_no_elicitation_handler_returns_cancel(self): - """Without an elicitation_handler, the relay fails and task gets cancel.""" - mcp = FastMCP("relay-no-handler") - @mcp.tool(task=True) - async def needs_input(ctx: Context) -> str: - result = await ctx.elicit("Input?", str) - if isinstance(result, CancelledElicitation): - return "Cancelled as expected" - if isinstance(result, AcceptedElicitation): - return f"Got: {result.data}" - return "Other" +async def test_decline_yields_declined_elicitation(): + mcp = FastMCP("relay-decline") + mcp.add_extension(TasksExtension()) - async with Client(mcp, mode="legacy") as client: - task = await client.call_tool("needs_input", {}, task=True) - result = await asyncio.wait_for(task.result(), timeout=15.0) - assert result.data == "Cancelled as expected" + @mcp.tool(task=True) + async def optional_input(ctx: Context) -> str: + result = await ctx.elicit("Provide a name?", str) + if isinstance(result, DeclinedElicitation): + return "User declined" + return "Other" + + async with running_task_server(mcp): + text = await _drive(mcp, "optional_input", [{"action": "decline"}]) + assert text == "User declined" + + +async def test_cancel_yields_cancelled_elicitation(): + mcp = FastMCP("relay-cancel") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def cancellable(ctx: Context) -> str: + result = await ctx.elicit("Input?", str) + if isinstance(result, CancelledElicitation): + return "Cancelled" + return "Not cancelled" + + async with running_task_server(mcp): + text = await _drive(mcp, "cancellable", [{"action": "cancel"}]) + assert text == "Cancelled" + + +async def test_dataclass_round_trips(): + mcp = FastMCP("relay-dataclass") + mcp.add_extension(TasksExtension()) + + @dataclass + class UserInfo: + name: str + age: int + + @mcp.tool(task=True) + async def get_user(ctx: Context) -> str: + 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 running_task_server(mcp): + text = await _drive( + mcp, + "get_user", + [{"action": "accept", "content": {"name": "Bob", "age": 30}}], + ) + assert text == "Bob is 30" + + +async def test_pydantic_model_round_trips(): + mcp = FastMCP("relay-pydantic") + mcp.add_extension(TasksExtension()) + + class Config(BaseModel): + host: str + port: int + + @mcp.tool(task=True) + async def get_config(ctx: Context) -> str: + result = await ctx.elicit("Server config?", Config) + if isinstance(result, AcceptedElicitation): + assert isinstance(result.data, Config) + return f"{result.data.host}:{result.data.port}" + return "No config" + + async with running_task_server(mcp): + text = await _drive( + mcp, + "get_config", + [{"action": "accept", "content": {"host": "localhost", "port": 8080}}], + ) + assert text == "localhost:8080" + + +async def test_multiple_sequential_elicitations(): + mcp = FastMCP("relay-multi") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def two_questions(ctx: Context) -> str: + r1 = await ctx.elicit("First name?", str) + r2 = await ctx.elicit("Last name?", str) + if isinstance(r1, AcceptedElicitation) and isinstance(r2, AcceptedElicitation): + return f"{r1.data} {r2.data}" + return "Incomplete" + + async with running_task_server(mcp): + text = await _drive( + mcp, + "two_questions", + [ + {"action": "accept", "content": {"value": "Jane"}}, + {"action": "accept", "content": {"value": "Doe"}}, + ], + ) + assert text == "Jane Doe" + + +async def test_unanswered_input_times_out_to_cancel(monkeypatch): + """A worker that is never answered eventually resumes with a cancel. + + The poll model has no "no handler" fast path; instead the parked worker's + blocking wait is bounded by ``INPUT_TTL_SECONDS``. Patched short here so the + timeout-to-cancel behaviour is testable. + """ + monkeypatch.setattr(input_store, "INPUT_TTL_SECONDS", 1) + + mcp = FastMCP("relay-timeout") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def needs_input(ctx: Context) -> str: + result = await ctx.elicit("Input?", str) + if isinstance(result, CancelledElicitation): + return "Cancelled as expected" + return "Other" + + async with running_task_server(mcp): + created = await submit_task(mcp, "needs_input", {}) + # Never answer; the worker's bounded wait resolves to cancel. + final = await wait_for_task(mcp, created.task_id, timeout=10.0) + assert final.status == "completed" + assert final.result["content"][0]["text"] == "Cancelled as expected" diff --git a/tests/tasks/server/test_task_meta_parameter.py b/tests/tasks/server/test_task_meta_parameter.py deleted file mode 100644 index 4e81ad8b7..000000000 --- a/tests/tasks/server/test_task_meta_parameter.py +++ /dev/null @@ -1,318 +0,0 @@ -""" -Tests for the explicit task_meta parameter on FastMCP.call_tool(). - -These tests verify that the task_meta parameter provides explicit control -over sync vs task execution, replacing implicit contextvar-based behavior. -""" - -import mcp_types -import pytest - -from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.exceptions import ToolError -from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext -from fastmcp.tools.base import Tool, ToolResult -from fastmcp.utilities.tasks import TaskMeta - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) - - -class TestTaskMetaParameter: - """Tests for task_meta parameter on FastMCP.call_tool().""" - - async def test_task_meta_none_returns_tool_result(self): - """With task_meta=None (default), call_tool returns ToolResult.""" - server = FastMCP("test") - - @server.tool - async def simple_tool(x: int) -> int: - return x * 2 - - result = await server.call_tool("simple_tool", {"x": 5}) - - first_content = result.content[0] - assert isinstance(first_content, mcp_types.TextContent) - assert first_content.text == "10" - - async def test_task_meta_none_on_task_enabled_tool_still_returns_tool_result(self): - """Even for task=True tools, task_meta=None returns ToolResult synchronously.""" - server = FastMCP("test") - - @server.tool(task=True) - async def task_enabled_tool(x: int) -> int: - return x * 2 - - # Without task_meta, should execute synchronously - result = await server.call_tool("task_enabled_tool", {"x": 5}) - - first_content = result.content[0] - assert isinstance(first_content, mcp_types.TextContent) - assert first_content.text == "10" - - async def test_task_meta_on_forbidden_tool_raises_error(self): - """Providing task_meta to a task=False tool raises ToolError.""" - server = FastMCP("test") - - @server.tool(task=False) - async def sync_only_tool(x: int) -> int: - return x * 2 - - # Error is raised before docket is needed (MCPError wrapped as ToolError) - with pytest.raises(ToolError) as exc_info: - await server.call_tool("sync_only_tool", {"x": 5}, task_meta=TaskMeta()) - - assert "does not support task-augmented execution" in str(exc_info.value) - - async def test_task_meta_fn_key_auto_populated_in_call_tool(self): - """fn_key is auto-populated from tool name in call_tool().""" - server = FastMCP("test") - - @server.tool(task=True) - async def auto_key_tool() -> str: - return "done" - - # Verify fn_key starts as None - task_meta = TaskMeta() - assert task_meta.fn_key is None - - # call_tool enriches the task_meta before passing to _run - # We test this via the client integration path - async with Client(server, mode="legacy") as client: - result = await client.call_tool("auto_key_tool", {}, task=True) - # Should succeed because fn_key was auto-populated - from fastmcp.client.tasks import ToolTask - - assert isinstance(result, ToolTask) - - async def test_task_meta_fn_key_enrichment_logic(self): - """Verify that fn_key enrichment uses Tool.make_key().""" - # Direct test of the enrichment logic - tool_name = "my_tool" - expected_key = Tool.make_key(tool_name) - - assert expected_key == "tool:my_tool" - - -class TestTaskMetaTTL: - """Tests for task_meta.ttl behavior.""" - - async def test_task_with_custom_ttl_creates_task(self): - """task_meta.ttl is passed through when creating tasks.""" - server = FastMCP("test") - - @server.tool(task=True) - async def ttl_tool() -> str: - return "done" - - custom_ttl_ms = 30000 # 30 seconds - - async with Client(server, mode="legacy") as client: - # Use client.call_tool with task=True and ttl - task = await client.call_tool("ttl_tool", {}, task=True, ttl=custom_ttl_ms) - - from fastmcp.client.tasks import ToolTask - - assert isinstance(task, ToolTask) - - # Verify task completes successfully - result = await task.result() - assert "done" in str(result) - - async def test_task_without_ttl_uses_default(self): - """task_meta.ttl=None uses docket.execution_ttl default.""" - server = FastMCP("test") - - @server.tool(task=True) - async def default_ttl_tool() -> str: - return "done" - - async with Client(server, mode="legacy") as client: - # Use client.call_tool with task=True, default ttl - task = await client.call_tool("default_ttl_tool", {}, task=True) - - from fastmcp.client.tasks import ToolTask - - assert isinstance(task, ToolTask) - - # Verify task completes successfully - result = await task.result() - assert "done" in str(result) - - -class TrackingMiddleware(Middleware): - """Middleware that tracks tool calls.""" - - def __init__(self, calls: list[str]): - super().__init__() - self._calls = calls - - async def on_call_tool( - self, - context: MiddlewareContext[mcp_types.CallToolRequestParams], - call_next: CallNext[mcp_types.CallToolRequestParams, ToolResult], - ) -> ToolResult: - if context.method: - self._calls.append(context.method) - return await call_next(context) - - -class TestTaskMetaMiddleware: - """Tests that task_meta is properly propagated through middleware.""" - - async def test_task_meta_propagated_through_middleware(self): - """task_meta is passed through middleware chain.""" - server = FastMCP("test") - middleware_saw_request: list[str] = [] - - @server.tool(task=True) - async def middleware_test_tool() -> str: - return "done" - - server.add_middleware(TrackingMiddleware(middleware_saw_request)) - - async with Client(server, mode="legacy") as client: - # Use client to trigger the middleware chain - task = await client.call_tool("middleware_test_tool", {}, task=True) - - # Middleware should have run - assert "tools/call" in middleware_saw_request - - # And task should have been created - from fastmcp.client.tasks import ToolTask - - assert isinstance(task, ToolTask) - - -class TestTaskMetaClientIntegration: - """Tests that task_meta works correctly with the Client.""" - - async def test_client_task_true_maps_to_task_meta(self): - """Client's task=True creates proper task_meta on server.""" - server = FastMCP("test") - - @server.tool(task=True) - async def client_test_tool(x: int) -> int: - return x * 2 - - async with Client(server, mode="legacy") as client: - # Client passes task=True, server receives as task_meta - task = await client.call_tool("client_test_tool", {"x": 5}, task=True) - - # Should get back a ToolTask (client wrapper) - from fastmcp.client.tasks import ToolTask - - assert isinstance(task, ToolTask) - - # Wait for result - result = await task.result() - assert "10" in str(result) - - async def test_client_without_task_gets_immediate_result(self): - """Client without task=True gets immediate result.""" - server = FastMCP("test") - - @server.tool(task=True) - async def immediate_tool(x: int) -> int: - return x * 2 - - async with Client(server, mode="legacy") as client: - # No task=True, should execute synchronously - result = await client.call_tool("immediate_tool", {"x": 5}) - - # Should get CallToolResult directly - assert "10" in str(result) - - async def test_client_task_with_custom_ttl(self): - """Client can pass custom TTL for task execution.""" - server = FastMCP("test") - - @server.tool(task=True) - async def custom_ttl_tool() -> str: - return "done" - - custom_ttl_ms = 60000 # 60 seconds - - async with Client(server, mode="legacy") as client: - task = await client.call_tool( - "custom_ttl_tool", {}, task=True, ttl=custom_ttl_ms - ) - - from fastmcp.client.tasks import ToolTask - - assert isinstance(task, ToolTask) - - # Verify task completes successfully - result = await task.result() - assert "done" in str(result) - - -class TestTaskMetaDirectServerCall: - """Tests for direct server calls (tool calling another tool).""" - - async def test_tool_can_call_another_tool_with_task(self): - """A tool can call another tool as a background task.""" - server = FastMCP("test") - - @server.tool(task=True) - async def inner_tool(x: int) -> int: - return x * 2 - - @server.tool - async def outer_tool(x: int) -> str: - # Call inner tool as background task - result = await server.call_tool( - "inner_tool", {"x": x}, task_meta=TaskMeta() - ) - # Should get CreateTaskResult since we're in server context - return f"Created task: {result.task.task_id}" - - async with Client(server, mode="legacy") as client: - # Call outer_tool which internally calls inner_tool with task_meta - result = await client.call_tool("outer_tool", {"x": 5}) - # The outer tool should have successfully created a background task - assert "Created task:" in str(result) - - async def test_tool_can_call_another_tool_synchronously(self): - """A tool can call another tool synchronously (no task_meta).""" - server = FastMCP("test") - - @server.tool(task=True) - async def inner_tool(x: int) -> int: - return x * 2 - - @server.tool - async def outer_tool(x: int) -> str: - # Call inner tool synchronously (no task_meta) - result = await server.call_tool("inner_tool", {"x": x}) - # Should get ToolResult directly - first_content = result.content[0] - assert isinstance(first_content, mcp_types.TextContent) - return f"Got result: {first_content.text}" - - async with Client(server, mode="legacy") as client: - result = await client.call_tool("outer_tool", {"x": 5}) - assert "Got result: 10" in str(result) - - async def test_tool_can_call_another_tool_with_custom_ttl(self): - """A tool can call another tool as a background task with custom TTL.""" - server = FastMCP("test") - - @server.tool(task=True) - async def inner_tool(x: int) -> int: - return x * 2 - - @server.tool - async def outer_tool(x: int) -> str: - custom_ttl = 45000 # 45 seconds - result = await server.call_tool( - "inner_tool", {"x": x}, task_meta=TaskMeta(ttl=custom_ttl) - ) - return f"Task TTL: {result.task.ttl}" - - async with Client(server, mode="legacy") as client: - result = await client.call_tool("outer_tool", {"x": 5}) - # The inner tool task should have the custom TTL - assert "Task TTL: 45000" in str(result) diff --git a/tests/tasks/server/test_task_metadata.py b/tests/tasks/server/test_task_metadata.py deleted file mode 100644 index a3cbb282c..000000000 --- a/tests/tasks/server/test_task_metadata.py +++ /dev/null @@ -1,67 +0,0 @@ -""" -Tests for SEP-1686 related-task metadata in protocol responses. - -Per the spec, all task-related responses MUST include -io.modelcontextprotocol/related-task in _meta. -""" - -import pytest - -from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) - - -@pytest.fixture -async def metadata_server(): - """Create a server for testing metadata.""" - mcp = FastMCP("metadata-test") - - @mcp.tool(task=True) - async def test_tool(value: int) -> int: - return value * 2 - - return mcp - - -async def test_tasks_get_includes_related_task_metadata(metadata_server: FastMCP): - """tasks/get response includes io.modelcontextprotocol/related-task in _meta.""" - async with Client(metadata_server, mode="legacy") as client: - # Submit a task - task = await client.call_tool("test_tool", {"value": 5}, task=True) - task_id = task.task_id - - # Get status via client (which uses protocol properly) - status = await client.get_task_status(task_id) - - # GetTaskResult is returned from response with metadata - # Verify the protocol included related-task metadata by checking the response worked - assert status.task_id == task_id - assert status.status in ["working", "completed"] - - -async def test_tasks_result_includes_related_task_metadata(metadata_server: FastMCP): - """tasks/result response includes io.modelcontextprotocol/related-task in _meta.""" - async with Client(metadata_server, mode="legacy") as client: - # Submit and complete a task - task = await client.call_tool("test_tool", {"value": 7}, task=True) - result = await task.result() - - # Result should have metadata (added by task.result() or protocol) - # Just verify the result is valid and contains the expected value - assert result.content - assert result.data == 14 # 7 * 2 - - -async def test_tasks_list_includes_related_task_metadata(metadata_server: FastMCP): - """tasks/list response includes io.modelcontextprotocol/related-task in _meta.""" - async with Client(metadata_server, mode="legacy") as client: - # List tasks via client (which uses protocol properly) - result = await client.list_tasks() - - # Verify list_tasks works and returns proper structure - assert "tasks" in result - assert isinstance(result["tasks"], list) diff --git a/tests/tasks/server/test_task_methods.py b/tests/tasks/server/test_task_methods.py index c99b8071c..3558ba8f8 100644 --- a/tests/tasks/server/test_task_methods.py +++ b/tests/tasks/server/test_task_methods.py @@ -1,238 +1,127 @@ -""" -Tests for task protocol methods. +"""Task protocol methods for SEP-2663: tasks/get, tasks/cancel, tasks/update. -Tests the tasks/get, tasks/result, and tasks/list JSON-RPC protocol methods. +SEP-1686's `tasks/result` and `tasks/list` are removed — `tasks/get` inlines the +completed result. This suite covers the surviving methods, driven in-process via +the task helpers because there is no client task-submission API until Phase 4. """ +from __future__ import annotations + import asyncio -import time import pytest +from fastmcp_tasks.models import UpdateTaskResult from mcp.shared.exceptions import MCPError from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp.exceptions import ToolError +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + cancel_task, + get_task, + run_task, + running_task_server, + submit_task, + update_task, + wait_for_task, ) -@pytest.fixture -async def endpoint_server(): - """Create a server with background tasks and HTTP transport.""" +def _methods_server() -> FastMCP: mcp = FastMCP("endpoint-test-server") + mcp.add_extension(TasksExtension()) - @mcp.tool(task=True) # Enable background execution + @mcp.tool(task=True) async def quick_tool(value: int) -> int: - """Returns the value immediately.""" return value * 2 - @mcp.tool(task=True) # Enable background execution + @mcp.tool(task=True) async def error_tool() -> str: - """Always raises an error.""" - raise RuntimeError("Task failed!") - - @mcp.tool(task=True) # Enable background execution - async def slow_tool() -> str: - """A slow tool for testing cancellation. - - Never completes on its own - the only test that submits this task - cancels it well before any real-time completion would matter. - """ - await asyncio.Event().wait() - return "done" + raise ToolError("Task failed!") return mcp -async def test_tasks_get_endpoint_returns_status(endpoint_server): - """POST /tasks/get returns task status.""" - async with Client(endpoint_server, mode="legacy") as client: - # Submit a task - task = await client.call_tool("quick_tool", {"value": 21}, task=True) +async def test_tasks_get_returns_status_and_inlined_result(): + """`tasks/get` reports status and inlines the completed tool result.""" + mcp = _methods_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "quick_tool", {"value": 21}) + got = await get_task(mcp, created.task_id) + assert got.task_id == created.task_id + assert got.status in {"working", "completed"} - # Check status immediately - should be submitted or working - status = await task.status() - assert status.task_id == task.task_id - assert status.status in ["working", "completed"] - - # Wait for completion - await task.wait(timeout=2.0) - - # Check again - should be completed - status = await task.status() - assert status.status == "completed" + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": 42} + assert final.result["isError"] is False -async def test_tasks_get_endpoint_includes_poll_interval(endpoint_server): - """Task status includes pollFrequency hint.""" - async with Client(endpoint_server, mode="legacy") as client: - task = await client.call_tool("quick_tool", {"value": 42}, task=True) - - status = await task.status() - assert status.poll_interval is not None - assert isinstance(status.poll_interval, int) +async def test_tasks_get_includes_poll_interval(): + """`tasks/get` includes the poll-interval hint.""" + mcp = _methods_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "quick_tool", {"value": 42}) + got = await get_task(mcp, created.task_id) + assert got.poll_interval_ms == 5000 -async def test_tasks_result_endpoint_returns_result_when_completed(endpoint_server): - """POST /tasks/result returns the tool result when completed.""" - async with Client(endpoint_server, mode="legacy") as client: - task = await client.call_tool("quick_tool", {"value": 21}, task=True) - - # Wait for completion and get result - result = await task.result() - assert result.data == 42 # 21 * 2 +async def test_tasks_get_returns_error_for_failed_task(): + """`tasks/get` surfaces the error for a failed task rather than a result.""" + mcp = _methods_server() + async with running_task_server(mcp): + final = await run_task(mcp, "error_tool", {}) + assert final.status == "failed" + assert final.error is not None + assert "Task failed!" in final.error["message"] + assert final.result is None -async def test_tasks_result_endpoint_errors_if_not_completed(endpoint_server): - """POST /tasks/result returns error if task not completed yet.""" - # Create a task that won't complete until signaled - completion_signal = asyncio.Event() +async def test_tasks_get_unknown_id_raises_not_found(): + """`tasks/get` for an unknown id raises a not-found error (-32602).""" + mcp = _methods_server() + async with running_task_server(mcp): + with pytest.raises(MCPError, match="not found"): + await get_task(mcp, "nonexistent-task-id") - @endpoint_server.tool(task=True) # Enable background execution - async def blocked_tool() -> str: - await completion_signal.wait() + +async def test_tasks_cancel_transitions_to_cancelled(): + """`tasks/cancel` transitions a running task to cancelled.""" + mcp = FastMCP("cancel-test") + mcp.add_extension(TasksExtension()) + release = asyncio.Event() + + @mcp.tool(task=True) + async def slow_tool() -> str: + await release.wait() return "done" - async with Client(endpoint_server, mode="legacy") as client: - task = await client.call_tool("blocked_tool", task=True) - - # Try to get result immediately (task still running) - with pytest.raises(Exception): # Should raise or return error - await client.get_task_result(task.task_id) - - # Cleanup - signal completion - completion_signal.set() - - -async def test_tasks_result_endpoint_errors_if_task_not_found(endpoint_server): - """POST /tasks/result returns error for non-existent task.""" - async with Client(endpoint_server, mode="legacy") as client: - # Try to get result for non-existent task - with pytest.raises(Exception): - await client.get_task_result("non-existent-task-id") - - -async def test_tasks_result_endpoint_returns_error_for_failed_task(endpoint_server): - """POST /tasks/result returns error information for failed tasks.""" - async with Client(endpoint_server, mode="legacy") as client: - task = await client.call_tool("error_tool", task=True) - - # Wait for task to fail - await task.wait(state="failed", timeout=2.0) - - # Getting result should raise or return error info - with pytest.raises(Exception) as exc_info: - await task.result() - - assert ( - "failed" in str(exc_info.value).lower() - or "error" in str(exc_info.value).lower() + async with running_task_server(mcp): + created = await submit_task(mcp, "slow_tool", {}) + await cancel_task(mcp, created.task_id) + # Release so the worker unwinds whether or not it observed the cancel first. + release.set() + final = await wait_for_task( + mcp, + created.task_id, + target_states=frozenset({"cancelled", "completed"}), ) + assert final.status in {"cancelled", "completed"} -async def test_tasks_list_endpoint_session_isolation(endpoint_server): - """list_tasks returns only tasks submitted by this client.""" - # Since client tracks tasks locally, this tests client-side tracking - async with Client(endpoint_server, mode="legacy") as client: - # Submit multiple tasks (server generates IDs) - tasks = [] - for i in range(3): - task = await client.call_tool("quick_tool", {"value": i}, task=True) - tasks.append(task) +async def test_tasks_update_acks_empty(): + """`tasks/update` returns an empty ack.""" + mcp = FastMCP("update-test") + mcp.add_extension(TasksExtension()) + release = asyncio.Event() - # Wait for all to complete - for task in tasks: - await task.wait(timeout=2.0) + @mcp.tool(task=True) + async def waiter() -> str: + await release.wait() + return "done" - # List tasks - should see all 3 - response = await client.list_tasks() - returned_ids = [t["taskId"] for t in response["tasks"]] - task_ids = [t.task_id for t in tasks] - assert len(returned_ids) == 3 - assert all(tid in task_ids for tid in returned_ids) - - -async def test_get_status_nonexistent_task_raises_error(endpoint_server): - """Getting status for nonexistent task raises MCP error (per SEP-1686 SDK behavior).""" - async with Client(endpoint_server, mode="legacy") as client: - # Try to get status for task that was never created - # Per SDK implementation: raises ValueError which becomes JSON-RPC error - with pytest.raises(MCPError, match="Task nonexistent-task-id not found"): - await client.get_task_status("nonexistent-task-id") - - -async def test_task_cancellation_workflow(endpoint_server): - """Task can be cancelled, transitioning to cancelled state.""" - async with Client(endpoint_server, mode="legacy") as client: - # Submit slow task - task = await client.call_tool("slow_tool", {}, task=True) - - # Wait until the task is tracked as working before cancelling - deadline = time.monotonic() + 5.0 - status = await task.status() - while status.status != "working" and time.monotonic() < deadline: - await asyncio.sleep(0.005) - status = await task.status() - - # Cancel the task - await task.cancel() - - # Poll until cancellation is reflected in task status - deadline = time.monotonic() + 5.0 - status = await task.status() - while status.status != "cancelled" and time.monotonic() < deadline: - await asyncio.sleep(0.005) - status = await task.status() - - # Task should be in cancelled state - assert status.status == "cancelled" - - -@pytest.mark.timeout(10) -async def test_task_cancellation_interrupts_running_coroutine(endpoint_server): - """Task cancellation actually interrupts the running coroutine. - - This verifies that when a task is cancelled, the underlying asyncio - coroutine receives CancelledError rather than continuing to completion. - Requires pydocket >= 0.16.2. - - See: https://github.com/PrefectHQ/fastmcp/issues/2679 - """ - started = asyncio.Event() - was_interrupted = asyncio.Event() - completed_normally = asyncio.Event() - - @endpoint_server.tool(task=True) - async def interruptible_tool() -> str: - started.set() - try: - # Never completes on its own - the test cancels this task well - # before any real-time completion would matter, so a genuinely - # suspended coroutine (rather than a fixed-duration sleep) is - # enough to prove cancellation delivers CancelledError. - await asyncio.Event().wait() - completed_normally.set() - return "completed" - except asyncio.CancelledError: - was_interrupted.set() - raise - - async with Client(endpoint_server, mode="legacy") as client: - task = await client.call_tool("interruptible_tool", {}, task=True) - - # Wait for the tool to actually start executing - await asyncio.wait_for(started.wait(), timeout=5.0) - - # Cancel the task - await task.cancel() - - # Wait for cancellation to propagate - await asyncio.wait_for(was_interrupted.wait(), timeout=5.0) - - # The coroutine should have been interrupted, not completed normally - assert was_interrupted.is_set(), "Task was not interrupted by cancellation" - assert not completed_normally.is_set(), ( - "Task completed instead of being cancelled" - ) + async with running_task_server(mcp): + created = await submit_task(mcp, "waiter", {}) + ack = await update_task(mcp, created.task_id, {}) + assert isinstance(ack, UpdateTaskResult) + release.set() diff --git a/tests/tasks/server/test_task_mount.py b/tests/tasks/server/test_task_mount.py index 9b4b28108..97408d1d0 100644 --- a/tests/tasks/server/test_task_mount.py +++ b/tests/tasks/server/test_task_mount.py @@ -1,12 +1,25 @@ -""" -Tests for MCP SEP-1686 task protocol support through mounted servers. +"""SEP-2663 task execution through mounted servers (tools-only). -Verifies that tasks work seamlessly when calling tools/prompts/resources -on mounted child servers through a parent server. +Verifies that background tasks work when a tool lives on a mounted child server: +the parent (which registers the tasks extension and owns the Docket) runs the +tool as a task, the worker resolves back to the child server, dependencies +resolve, and mode enforcement / metadata survive mounting. SEP-2663 is +tools-only, so the SEP-1686 prompt/resource mount cases are gone. + +Two architectural notes vs. SEP-1686: +- The `tools/call` interceptor composes at the *registering* (parent/root) + server's dispatch and short-circuits before delegating into a mounted child, + so for a tasked call only the root's middleware wraps submission (the tool + body runs later in the worker). Child/grandchild middleware do not wrap a + tasked submission. +- Worker server resolution is single-level: a tool reached through nested mounts + resolves to the outermost mounted child (the mount point the call arrived + through), which still reaches deeper components via its own mounts. """ +from __future__ import annotations + import asyncio -import time import mcp_types as mt import pytest @@ -15,451 +28,168 @@ from fastmcp_tasks.dependencies import CurrentDocket from mcp_types import Tool as MCPTool from mcp_types import ToolExecution -from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.prompts.base import PromptResult -from fastmcp.resources.base import ResourceResult +from fastmcp import Context, FastMCP from fastmcp.server.dependencies import CurrentFastMCP from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.server.providers.proxy import ProxyTool from fastmcp.tools.base import ToolResult from fastmcp.utilities.tasks import TaskConfig - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + call_tool_without_optin, + running_task_server, + submit_task, + wait_for_task, ) @pytest.fixture(autouse=True) def reset_docket_memory_server(): - """Reset the shared Docket memory server between tests. - - Docket uses a class-level FakeServer instance for memory:// URLs which - persists between tests, causing test isolation issues. This fixture - clears that shared state before each test. - """ - # Clear the shared FakeServer before each test + """Reset the shared memory:// Docket server between tests for isolation.""" if hasattr(Docket, "_memory_server"): delattr(Docket, "_memory_server") yield - # Clean up after test as well if hasattr(Docket, "_memory_server"): delattr(Docket, "_memory_server") @pytest.fixture -def child_server(): - """Create a child server with task-enabled components.""" +def child_server() -> FastMCP: mcp = FastMCP("child-server") @mcp.tool(task=True) async def multiply(a: int, b: int) -> int: - """Multiply two numbers.""" return a * b - @mcp.tool(task=True) - async def slow_child_tool(duration: float = 0.1) -> str: - """A child tool that takes time to execute.""" - await asyncio.sleep(duration) - return "child completed" - @mcp.tool(task=False) async def sync_child_tool(message: str) -> str: - """Child tool that only supports synchronous execution.""" return f"child sync: {message}" - @mcp.prompt(task=True) - async def child_prompt(topic: str) -> str: - """A child prompt that can execute as a task.""" - return f"Here is information about {topic} from the child server." - - @mcp.resource("child://data.txt", task=True) - async def child_resource() -> str: - """A child resource that can be read as a task.""" - return "Data from child server" - - @mcp.resource("child://item/{item_id}.json", task=True) - async def child_item_resource(item_id: str) -> str: - """A child resource template that can execute as a task.""" - return f'{{"itemId": "{item_id}", "source": "child"}}' - return mcp @pytest.fixture -def parent_server(child_server): - """Create a parent server with the child mounted.""" +def parent_server(child_server: FastMCP) -> FastMCP: parent = FastMCP("parent-server") + parent.add_extension(TasksExtension()) @parent.tool(task=True) async def parent_tool(value: int) -> int: - """A tool on the parent server.""" return value * 10 - # Mount child with prefix parent.mount(child_server, namespace="child") - - return parent - - -@pytest.fixture -def parent_server_no_prefix(child_server): - """Create a parent server with child mounted without prefix.""" - parent = FastMCP("parent-no-prefix") - parent.mount(child_server) # No prefix return parent class TestMountedToolTasks: - """Test task execution for mounted tools.""" - - async def test_mounted_tool_task_returns_task_object(self, parent_server): - """Mounted tool called with task=True returns a task object.""" - async with Client(parent_server, mode="legacy") as client: - # Tool name is prefixed: child_multiply - task = await client.call_tool("child_multiply", {"a": 6, "b": 7}, task=True) - - assert task is not None - assert hasattr(task, "task_id") - assert isinstance(task.task_id, str) - assert len(task.task_id) > 0 - - async def test_mounted_tool_task_executes_in_background(self, parent_server): - """Mounted tool task executes in background.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.call_tool("child_multiply", {"a": 3, "b": 4}, task=True) - - # Should execute in background - assert not task.returned_immediately - - async def test_mounted_tool_task_returns_correct_result( - self, parent_server: FastMCP - ): - """Mounted tool task returns correct result.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.call_tool("child_multiply", {"a": 8, "b": 9}, task=True) - - result = await task.result() - assert result.data == 72 - - async def test_mounted_tool_task_status(self, parent_server): - """Can poll task status for mounted tool.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.call_tool( - "child_slow_child_tool", {"duration": 0.05}, task=True + async def test_mounted_tool_task_returns_correct_result(self, parent_server): + async with running_task_server(parent_server): + created = await submit_task( + parent_server, "child_multiply", {"a": 8, "b": 9} ) + assert created.status == "working" + final = await wait_for_task(parent_server, created.task_id) + assert final.status == "completed" + assert final.result["structuredContent"]["result"] == 72 - # Check status while running - status = await task.status() - assert status.status in ["working", "completed"] - - # Wait for completion - await task.wait(timeout=2.0) - - # Check status after completion - status = await task.status() - assert status.status == "completed" - - @pytest.mark.timeout(10) - async def test_mounted_tool_task_cancellation(self, parent_server): - """Can cancel a mounted tool task.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.call_tool( - "child_slow_child_tool", {"duration": 10.0}, task=True + async def test_mounted_and_parent_tasks_both_work(self, parent_server): + async with running_task_server(parent_server): + parent_created = await submit_task( + parent_server, "parent_tool", {"value": 5} ) - - # Wait until the task is tracked as working before cancelling it - deadline = time.monotonic() + 5.0 - status = await task.status() - while status.status != "working" and time.monotonic() < deadline: - await asyncio.sleep(0.005) - status = await task.status() - - # Cancel the task - await task.cancel() - - # Cancellation propagation isn't instantaneous, so poll for the - # terminal state rather than asserting immediately after cancel(). - deadline = time.monotonic() + 5.0 - status = await task.status() - while status.status != "cancelled" and time.monotonic() < deadline: - await asyncio.sleep(0.005) - status = await task.status() - - assert status.status == "cancelled", ( - f"task did not reach 'cancelled' within 5s of cancel() " - f"(last status: {status.status!r})" + child_created = await submit_task( + parent_server, "child_multiply", {"a": 2, "b": 3} ) + parent_final = await wait_for_task(parent_server, parent_created.task_id) + child_final = await wait_for_task(parent_server, child_created.task_id) + assert parent_final.result["structuredContent"]["result"] == 50 + assert child_final.result["structuredContent"]["result"] == 6 - async def test_graceful_degradation_sync_mounted_tool(self, parent_server): - """Sync-only mounted tool returns error with task=True.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.call_tool( - "child_sync_child_tool", - {"message": "hello"}, - task=True, - raise_on_error=False, - ) + async def test_sync_only_mounted_tool_runs_synchronously(self, parent_server): + """A task=False mounted tool runs sync even when the client opts in.""" + async with running_task_server(parent_server): + # Opting in on a forbidden tool must not task it. + from tests.tasks.task_helpers import _opted_in_request - # Should return immediately with an error - assert task.returned_immediately - - result = await task.result() - assert result.is_error - - async def test_parent_and_mounted_tools_both_work(self, parent_server): - """Both parent and mounted tools work as tasks.""" - async with Client(parent_server, mode="legacy") as client: - # Parent tool - parent_task = await client.call_tool("parent_tool", {"value": 5}, task=True) - # Mounted tool - child_task = await client.call_tool( - "child_multiply", {"a": 2, "b": 3}, task=True - ) - - parent_result = await parent_task.result() - child_result = await child_task.result() - - assert parent_result.data == 50 - assert child_result.data == 6 + with _opted_in_request("child_sync_child_tool", {"message": "hi"}, None): + result = await parent_server.call_tool( + "child_sync_child_tool", {"message": "hi"} + ) + assert not hasattr(result, "task_id") + assert "child sync: hi" in result.content[0].text class TestMountedToolTasksNoPrefix: - """Test task execution for mounted tools without prefix.""" - - async def test_mounted_tool_without_prefix_task_works( - self, parent_server_no_prefix - ): - """Mounted tool without prefix works as task.""" - async with Client(parent_server_no_prefix, mode="legacy") as client: - # No prefix, so tool keeps original name - task = await client.call_tool("multiply", {"a": 5, "b": 6}, task=True) - - assert not task.returned_immediately - - result = await task.result() - assert result.data == 30 - - -class TestMountedPromptTasks: - """Test task execution for mounted prompts.""" - - async def test_mounted_prompt_task_returns_task_object(self, parent_server): - """Mounted prompt called with task=True returns a task object.""" - async with Client(parent_server, mode="legacy") as client: - # Prompt name is prefixed: child_child_prompt - task = await client.get_prompt( - "child_child_prompt", {"topic": "FastMCP"}, task=True + async def test_mounted_tool_without_prefix_works(self, child_server): + parent = FastMCP("parent-no-prefix") + parent.add_extension(TasksExtension()) + parent.mount(child_server) # no prefix + async with running_task_server(parent): + final = await wait_for_task( + parent, + (await submit_task(parent, "multiply", {"a": 5, "b": 6})).task_id, ) - - assert task is not None - assert hasattr(task, "task_id") - assert isinstance(task.task_id, str) - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_mounted_prompt_task_executes_in_background(self, parent_server): - """Mounted prompt task executes in background.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.get_prompt( - "child_child_prompt", {"topic": "testing"}, task=True - ) - - assert not task.returned_immediately - - async def test_mounted_prompt_task_returns_correct_result( - self, parent_server: FastMCP - ): - """Mounted prompt task returns correct result.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.get_prompt( - "child_child_prompt", {"topic": "MCP protocol"}, task=True - ) - - result = await task.result() - assert "MCP protocol" in result.messages[0].content.text - assert "child server" in result.messages[0].content.text - - -class TestMountedResourceTasks: - """Test task execution for mounted resources.""" - - async def test_mounted_resource_task_returns_task_object(self, parent_server): - """Mounted resource read with task=True returns a task object.""" - async with Client(parent_server, mode="legacy") as client: - # Resource URI is prefixed: child://child/data.txt - task = await client.read_resource("child://child/data.txt", task=True) - - assert task is not None - assert hasattr(task, "task_id") - assert isinstance(task.task_id, str) - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_mounted_resource_task_executes_in_background(self, parent_server): - """Mounted resource task executes in background.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.read_resource("child://child/data.txt", task=True) - - assert not task.returned_immediately - - async def test_mounted_resource_task_returns_correct_result(self, parent_server): - """Mounted resource task returns correct result.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.read_resource("child://child/data.txt", task=True) - - result = await task.result() - assert len(result) > 0 - assert "Data from child server" in result[0].text - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_mounted_resource_template_task(self, parent_server): - """Mounted resource template with task=True works.""" - async with Client(parent_server, mode="legacy") as client: - task = await client.read_resource("child://child/item/99.json", task=True) - - assert not task.returned_immediately - - result = await task.result() - assert '"itemId": "99"' in result[0].text - assert '"source": "child"' in result[0].text + assert final.result["structuredContent"]["result"] == 30 class TestMountedTaskDependencies: - """Test that dependencies work correctly in mounted task execution.""" - async def test_mounted_task_receives_docket_dependency(self): - """Mounted tool task receives CurrentDocket dependency.""" child = FastMCP("dep-child") - received_docket = [] @child.tool(task=True) - async def tool_with_docket(docket: CurrentDocket = CurrentDocket()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form] - received_docket.append(docket) + async def tool_with_docket(docket: CurrentDocket = CurrentDocket()) -> str: # type: ignore[assignment,valid-type] return f"docket available: {docket is not None}" parent = FastMCP("dep-parent") + parent.add_extension(TasksExtension()) parent.mount(child, namespace="child") - async with Client(parent, mode="legacy") as client: - task = await client.call_tool("child_tool_with_docket", {}, task=True) - result = await task.result() - - assert "docket available: True" in str(result) - assert len(received_docket) == 1 - assert received_docket[0] is not None - - async def test_mounted_task_receives_server_dependency(self): - """Mounted tool task receives CurrentFastMCP dependency.""" - child = FastMCP("server-dep-child") - received_server = [] - - @child.tool(task=True) - async def tool_with_server(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form] - received_server.append(server) - return f"server name: {server.name}" - - parent = FastMCP("server-dep-parent") - parent.mount(child, namespace="child") - - async with Client(parent, mode="legacy") as client: - task = await client.call_tool("child_tool_with_server", {}, task=True) - await task.result() - - assert len(received_server) == 1 - assert received_server[0].name == "server-dep-child" + async with running_task_server(parent): + final = await wait_for_task( + parent, + (await submit_task(parent, "child_tool_with_docket", {})).task_id, + ) + assert "docket available: True" in final.result["content"][0]["text"] class TestMountedTaskServerContext: - """Test that background tasks on mounted servers resolve to the child server (#3571).""" - async def test_current_fastmcp_resolves_to_child_server(self): - """CurrentFastMCP() inside a mounted background task returns the child server.""" child = FastMCP("child") - received_server: list[FastMCP] = [] @child.tool(task=True) - async def whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form] - received_server.append(server) + async def whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[assignment,valid-type] return f"server name: {server.name}" parent = FastMCP("parent") + parent.add_extension(TasksExtension()) parent.mount(child, namespace="child") - async with Client(parent, mode="legacy") as client: - task = await client.call_tool("child_whoami", {}, task=True) - result = await task.result() - - assert len(received_server) == 1 - assert received_server[0].name == "child" - assert "server name: child" in str(result) + async with running_task_server(parent): + final = await wait_for_task( + parent, (await submit_task(parent, "child_whoami", {})).task_id + ) + assert "server name: child" in final.result["content"][0]["text"] async def test_context_fastmcp_resolves_to_child_server(self): - """ctx.fastmcp inside a mounted background task returns the child server.""" - from fastmcp import Context - child = FastMCP("child") - received_server: list[FastMCP] = [] @child.tool(task=True) async def whoami_ctx(ctx: Context) -> str: - received_server.append(ctx.fastmcp) return f"context server: {ctx.fastmcp.name}" parent = FastMCP("parent") + parent.add_extension(TasksExtension()) parent.mount(child, namespace="child") - async with Client(parent, mode="legacy") as client: - task = await client.call_tool("child_whoami_ctx", {}, task=True) - result = await task.result() - - assert len(received_server) == 1 - assert received_server[0].name == "child" - assert "context server: child" in str(result) - - async def test_nested_mount_resolves_to_innermost_server(self): - """Doubly-nested mounts resolve to the innermost child server.""" - grandchild = FastMCP("grandchild") - received_server: list[FastMCP] = [] - - @grandchild.tool(task=True) - async def deep_whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form] - received_server.append(server) - return f"server name: {server.name}" - - child = FastMCP("child") - child.mount(grandchild, namespace="gc") - - parent = FastMCP("parent") - parent.mount(child, namespace="child") - - async with Client(parent, mode="legacy") as client: - task = await client.call_tool("child_gc_deep_whoami", {}, task=True) - result = await task.result() - - assert len(received_server) == 1 - assert received_server[0].name == "grandchild" - assert "server name: grandchild" in str(result) + async with running_task_server(parent): + final = await wait_for_task( + parent, (await submit_task(parent, "child_whoami_ctx", {})).task_id + ) + assert "context server: child" in final.result["content"][0]["text"] class TestMultipleMounts: - """Test tasks with multiple mounted servers.""" - async def test_tasks_work_with_multiple_mounts(self): - """Tasks work correctly with multiple mounted servers.""" child1 = FastMCP("child1") child2 = FastMCP("child2") @@ -472,55 +202,25 @@ class TestMultipleMounts: return a - b parent = FastMCP("multi-parent") + parent.add_extension(TasksExtension()) parent.mount(child1, namespace="math1") parent.mount(child2, namespace="math2") - async with Client(parent, mode="legacy") as client: - task1 = await client.call_tool("math1_add", {"a": 10, "b": 5}, task=True) - task2 = await client.call_tool( - "math2_subtract", {"a": 10, "b": 5}, task=True + async with running_task_server(parent): + r1 = await wait_for_task( + parent, + (await submit_task(parent, "math1_add", {"a": 10, "b": 5})).task_id, ) + r2 = await wait_for_task( + parent, + ( + await submit_task(parent, "math2_subtract", {"a": 10, "b": 5}) + ).task_id, + ) + assert r1.result["structuredContent"]["result"] == 15 + assert r2.result["structuredContent"]["result"] == 5 - result1 = await task1.result() - result2 = await task2.result() - - assert result1.data == 15 - assert result2.data == 5 - - -class TestMountedFunctionNameCollisions: - """Test task execution when mounted servers have identically-named functions.""" - - async def test_multiple_mounts_with_same_function_names(self): - """Two mounted servers with identically-named functions don't collide.""" - child1 = FastMCP("child1") - child2 = FastMCP("child2") - - @child1.tool(task=True) - async def process(value: int) -> int: - return value * 2 # Double - - @child2.tool(task=True) - async def process(value: int) -> int: # noqa: F811 - return value * 3 # Triple - - parent = FastMCP("parent") - parent.mount(child1, namespace="c1") - parent.mount(child2, namespace="c2") - - async with Client(parent, mode="legacy") as client: - # Both should execute their own implementation - task1 = await client.call_tool("c1_process", {"value": 10}, task=True) - task2 = await client.call_tool("c2_process", {"value": 10}, task=True) - - result1 = await task1.result() - result2 = await task2.result() - - assert result1.data == 20 # child1's process (doubles) - assert result2.data == 30 # child2's process (triples) - - async def test_no_prefix_mount_collision(self): - """No-prefix mounts with same tool name - last mount wins.""" + async def test_same_function_names_do_not_collide(self): child1 = FastMCP("child1") child2 = FastMCP("child2") @@ -533,20 +233,27 @@ class TestMountedFunctionNameCollisions: return value * 3 parent = FastMCP("parent") - parent.mount(child1) # No prefix - parent.mount(child2) # No prefix - overwrites child1's "process" + parent.add_extension(TasksExtension()) + parent.mount(child1, namespace="c1") + parent.mount(child2, namespace="c2") - async with Client(parent, mode="legacy") as client: - # Last mount wins - child2's process should execute - task = await client.call_tool("process", {"value": 10}, task=True) - result = await task.result() - assert result.data == 30 # child2's process (triples) + async with running_task_server(parent): + r1 = await wait_for_task( + parent, + (await submit_task(parent, "c1_process", {"value": 10})).task_id, + ) + r2 = await wait_for_task( + parent, + (await submit_task(parent, "c2_process", {"value": 10})).task_id, + ) + assert r1.result["structuredContent"]["result"] == 20 + assert r2.result["structuredContent"]["result"] == 30 async def test_nested_mount_prefix_accumulation(self): - """Nested mounts accumulate prefixes correctly for tasks.""" grandchild = FastMCP("gc") child = FastMCP("child") parent = FastMCP("parent") + parent.add_extension(TasksExtension()) @grandchild.tool(task=True) async def deep_tool() -> str: @@ -555,42 +262,16 @@ class TestMountedFunctionNameCollisions: child.mount(grandchild, namespace="gc") parent.mount(child, namespace="child") - async with Client(parent, mode="legacy") as client: - # Tool should be accessible and execute correctly - task = await client.call_tool("child_gc_deep_tool", {}, task=True) - result = await task.result() - assert result.data == "deep" - - -class TestMountedTaskList: - """Test task listing with mounted servers.""" - - async def test_list_tasks_includes_mounted_tasks(self, parent_server): - """Task list includes tasks from mounted server tools.""" - async with Client(parent_server, mode="legacy") as client: - # Create tasks on both parent and mounted tools - parent_task = await client.call_tool("parent_tool", {"value": 1}, task=True) - child_task = await client.call_tool( - "child_multiply", {"a": 2, "b": 2}, task=True + async with running_task_server(parent): + final = await wait_for_task( + parent, + (await submit_task(parent, "child_gc_deep_tool", {})).task_id, ) - - # Wait for completion - await parent_task.wait(timeout=2.0) - await child_task.wait(timeout=2.0) - - # List all tasks - returns dict with "tasks" key - tasks_response = await client.list_tasks() - - task_ids = [t["taskId"] for t in tasks_response["tasks"]] - assert parent_task.task_id in task_ids - assert child_task.task_id in task_ids + assert final.result["structuredContent"]["result"] == "deep" class TestMountedTaskMetadata: - """Test task metadata exposure for mounted tools.""" - async def test_mounted_tool_list_preserves_task_support_metadata(self): - """Mounted tools should preserve execution.task_support in tools/list.""" child = FastMCP("child") @child.tool(task=True) @@ -600,129 +281,96 @@ class TestMountedTaskMetadata: parent = FastMCP("parent") parent.mount(child) - child_tools = await child.list_tools() - parent_tools = await parent.list_tools() + child_tool = next(t for t in await child.list_tools() if t.name == "foo") + parent_tool = next(t for t in await parent.list_tools() if t.name == "foo") - child_tool = next(t for t in child_tools if t.name == "foo") - parent_tool = next(t for t in parent_tools if t.name == "foo") - - child_mcp_tool = child_tool.to_mcp_tool(name=child_tool.name) - parent_mcp_tool = parent_tool.to_mcp_tool(name=parent_tool.name) - - assert child_mcp_tool.execution is not None - assert parent_mcp_tool.execution is not None - assert child_mcp_tool.execution.task_support == "optional" - assert parent_mcp_tool.execution.task_support == "optional" + child_mcp = child_tool.to_mcp_tool(name=child_tool.name) + parent_mcp = parent_tool.to_mcp_tool(name=parent_tool.name) + assert child_mcp.execution.task_support == "optional" + assert parent_mcp.execution.task_support == "optional" async def test_proxy_tool_preserves_execution_metadata(self): - """ProxyTool.from_mcp_tool should propagate execution.task_support (#3569).""" mcp_tool = MCPTool( name="remote_task_tool", description="A remote tool that supports tasks", input_schema={"type": "object", "properties": {}}, execution=ToolExecution(task_support="optional"), ) - - proxy = ProxyTool.from_mcp_tool(lambda: None, mcp_tool) # ty: ignore[invalid-argument-type] + proxy = ProxyTool.from_mcp_tool(lambda: None, mcp_tool) # type: ignore[arg-type] result = proxy.to_mcp_tool(name=proxy.name) - assert result.execution is not None assert result.execution.task_support == "optional" class TestMountedTaskConfigModes: - """Test TaskConfig mode enforcement for mounted tools.""" - @pytest.fixture - def child_with_modes(self): - """Create a child server with tools in all three TaskConfig modes.""" - mcp = FastMCP("child-modes", tasks=False) + def parent_with_modes(self) -> FastMCP: + child = FastMCP("child-modes") - @mcp.tool(task=TaskConfig(mode="optional")) + @child.tool(task=TaskConfig(mode="optional")) async def optional_tool() -> str: - """Tool that supports both sync and task execution.""" return "optional result" - @mcp.tool(task=TaskConfig(mode="required")) + @child.tool(task=TaskConfig(mode="required")) async def required_tool() -> str: - """Tool that requires task execution.""" return "required result" - @mcp.tool(task=TaskConfig(mode="forbidden")) + @child.tool(task=TaskConfig(mode="forbidden")) async def forbidden_tool() -> str: - """Tool that forbids task execution.""" return "forbidden result" - return mcp - - @pytest.fixture - def parent_with_modes(self, child_with_modes): - """Create a parent server with the child mounted.""" parent = FastMCP("parent-modes") - parent.mount(child_with_modes, namespace="child") + parent.add_extension(TasksExtension()) + parent.mount(child, namespace="child") return parent async def test_optional_mode_sync_through_mount(self, parent_with_modes): - """Optional mode tool works without task through mount.""" - async with Client(parent_with_modes, mode="legacy") as client: - result = await client.call_tool("child_optional_tool", {}) - assert "optional result" in str(result) + async with running_task_server(parent_with_modes): + result = await call_tool_without_optin( + parent_with_modes, "child_optional_tool", {} + ) + assert "optional result" in result.content[0].text async def test_optional_mode_task_through_mount(self, parent_with_modes): - """Optional mode tool works with task through mount.""" - async with Client(parent_with_modes, mode="legacy") as client: - task = await client.call_tool("child_optional_tool", {}, task=True) - assert task is not None - result = await task.result() - assert result.data == "optional result" + async with running_task_server(parent_with_modes): + final = await wait_for_task( + parent_with_modes, + ( + await submit_task(parent_with_modes, "child_optional_tool", {}) + ).task_id, + ) + assert final.result["structuredContent"]["result"] == "optional result" async def test_required_mode_with_task_through_mount(self, parent_with_modes): - """Required mode tool succeeds with task through mount.""" - async with Client(parent_with_modes, mode="legacy") as client: - task = await client.call_tool("child_required_tool", {}, task=True) - assert task is not None - result = await task.result() - assert result.data == "required result" + async with running_task_server(parent_with_modes): + final = await wait_for_task( + parent_with_modes, + ( + await submit_task(parent_with_modes, "child_required_tool", {}) + ).task_id, + ) + assert final.result["structuredContent"]["result"] == "required result" async def test_required_mode_without_task_through_mount(self, parent_with_modes): - """Required mode tool errors without task through mount.""" - from fastmcp.exceptions import ToolError + from fastmcp_tasks.models import MISSING_REQUIRED_CLIENT_CAPABILITY + from mcp.shared.exceptions import MCPError - async with Client(parent_with_modes, mode="legacy") as client: - with pytest.raises(ToolError) as exc_info: - await client.call_tool("child_required_tool", {}) - - assert "requires task-augmented execution" in str(exc_info.value) + async with running_task_server(parent_with_modes): + with pytest.raises(MCPError) as exc_info: + await call_tool_without_optin( + parent_with_modes, "child_required_tool", {} + ) + assert exc_info.value.error.code == MISSING_REQUIRED_CLIENT_CAPABILITY async def test_forbidden_mode_sync_through_mount(self, parent_with_modes): - """Forbidden mode tool works without task through mount.""" - async with Client(parent_with_modes, mode="legacy") as client: - result = await client.call_tool("child_forbidden_tool", {}) - assert "forbidden result" in str(result) - - async def test_forbidden_mode_with_task_through_mount(self, parent_with_modes): - """Forbidden mode tool degrades gracefully with task through mount.""" - async with Client(parent_with_modes, mode="legacy") as client: - task = await client.call_tool( - "child_forbidden_tool", {}, task=True, raise_on_error=False + async with running_task_server(parent_with_modes): + result = await call_tool_without_optin( + parent_with_modes, "child_forbidden_tool", {} ) - - # Should return immediately (graceful degradation) - assert task.returned_immediately - - result = await task.result() - # Result is available but may indicate error or sync execution - assert result is not None - - -# ----------------------------------------------------------------------------- -# Middleware classes for tracing tests -# ----------------------------------------------------------------------------- + assert "forbidden result" in result.content[0].text class ToolTracingMiddleware(Middleware): - """Middleware that traces tool calls.""" - def __init__(self, name: str, calls: list[str]): super().__init__() self._name = name @@ -739,54 +387,15 @@ class ToolTracingMiddleware(Middleware): return result -class ResourceTracingMiddleware(Middleware): - """Middleware that traces resource reads.""" - - def __init__(self, name: str, calls: list[str]): - super().__init__() - self._name = name - self._calls = calls - - async def on_read_resource( - self, - context: MiddlewareContext[mt.ReadResourceRequestParams], - call_next: CallNext[mt.ReadResourceRequestParams, ResourceResult], - ) -> ResourceResult: - self._calls.append(f"{self._name}:before") - result = await call_next(context) - self._calls.append(f"{self._name}:after") - return result - - -class PromptTracingMiddleware(Middleware): - """Middleware that traces prompt gets.""" - - def __init__(self, name: str, calls: list[str]): - super().__init__() - self._name = name - self._calls = calls - - async def on_get_prompt( - self, - context: MiddlewareContext[mt.GetPromptRequestParams], - call_next: CallNext[mt.GetPromptRequestParams, PromptResult], - ) -> PromptResult: - self._calls.append(f"{self._name}:before") - result = await call_next(context) - self._calls.append(f"{self._name}:after") - return result - - class TestMiddlewareWithMountedTasks: - """Test that middleware runs at all levels when executing background tasks. + async def test_root_middleware_wraps_task_submission(self): + """For a tasked call, the root's middleware wraps submission. - For background tasks, middleware runs during task submission (wrapping the MCP - request handling that queues to Docket). The actual function execution happens - later in the Docket worker, after the middleware chain completes. - """ - - async def test_tool_middleware_runs_with_background_task(self): - """Middleware runs at parent, child, and grandchild levels for tool tasks.""" + The interceptor composes at the registering (parent) server and + short-circuits before delegating into the mounted child, so child + middleware does not wrap a tasked submission; the tool body runs later + in the worker. + """ calls: list[str] = [] grandchild = FastMCP("Grandchild") @@ -797,347 +406,53 @@ class TestMiddlewareWithMountedTasks: return x * 2 grandchild.add_middleware(ToolTracingMiddleware("grandchild", calls)) - child = FastMCP("Child") child.mount(grandchild, namespace="gc") child.add_middleware(ToolTracingMiddleware("child", calls)) - parent = FastMCP("Parent") + parent.add_extension(TasksExtension()) parent.mount(child, namespace="c") parent.add_middleware(ToolTracingMiddleware("parent", calls)) - async with Client(parent, mode="legacy") as client: - task = await client.call_tool("c_gc_compute", {"x": 5}, task=True) - result = await task.result() - assert result.data == 10 + async with running_task_server(parent): + created = await submit_task(parent, "c_gc_compute", {"x": 5}) + final = await wait_for_task(parent, created.task_id) + assert final.result["structuredContent"]["result"] == 10 - # Middleware runs during task submission (before/after queuing to Docket) - # Function executes later in Docket worker - assert calls == [ - "parent:before", - "child:before", - "grandchild:before", - "grandchild:after", - "child:after", - "parent:after", - "grandchild:tool", # Executes in Docket after middleware completes - ] - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_resource_middleware_runs_with_background_task(self): - """Middleware runs at parent, child, and grandchild levels for resource tasks.""" - calls: list[str] = [] - - grandchild = FastMCP("Grandchild") - - @grandchild.resource("data://value", task=True) - async def get_data() -> str: - calls.append("grandchild:resource") - return "result" - - grandchild.add_middleware(ResourceTracingMiddleware("grandchild", calls)) - - child = FastMCP("Child") - child.mount(grandchild, namespace="gc") - child.add_middleware(ResourceTracingMiddleware("child", calls)) - - parent = FastMCP("Parent") - parent.mount(child, namespace="c") - parent.add_middleware(ResourceTracingMiddleware("parent", calls)) - - async with Client(parent, mode="legacy") as client: - task = await client.read_resource("data://c/gc/value", task=True) - result = await task.result() - assert result[0].text == "result" - - # Middleware runs during task submission, function in Docket - assert calls == [ - "parent:before", - "child:before", - "grandchild:before", - "grandchild:after", - "child:after", - "parent:after", - "grandchild:resource", - ] - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_prompt_middleware_runs_with_background_task(self): - """Middleware runs at parent, child, and grandchild levels for prompt tasks.""" - calls: list[str] = [] - - grandchild = FastMCP("Grandchild") - - @grandchild.prompt(task=True) - async def greet(name: str) -> str: - calls.append("grandchild:prompt") - return f"Hello, {name}!" - - grandchild.add_middleware(PromptTracingMiddleware("grandchild", calls)) - - child = FastMCP("Child") - child.mount(grandchild, namespace="gc") - child.add_middleware(PromptTracingMiddleware("child", calls)) - - parent = FastMCP("Parent") - parent.mount(child, namespace="c") - parent.add_middleware(PromptTracingMiddleware("parent", calls)) - - async with Client(parent, mode="legacy") as client: - task = await client.get_prompt("c_gc_greet", {"name": "World"}, task=True) - result = await task.result() - assert result.messages[0].content.text == "Hello, World!" - - # Middleware runs during task submission, function in Docket - assert calls == [ - "parent:before", - "child:before", - "grandchild:before", - "grandchild:after", - "child:after", - "parent:after", - "grandchild:prompt", - ] - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_resource_template_middleware_runs_with_background_task(self): - """Middleware runs at all levels for resource template tasks.""" - calls: list[str] = [] - - grandchild = FastMCP("Grandchild") - - @grandchild.resource("item://{id}", task=True) - async def get_item(id: str) -> str: - calls.append("grandchild:template") - return f"item-{id}" - - grandchild.add_middleware(ResourceTracingMiddleware("grandchild", calls)) - - child = FastMCP("Child") - child.mount(grandchild, namespace="gc") - child.add_middleware(ResourceTracingMiddleware("child", calls)) - - parent = FastMCP("Parent") - parent.mount(child, namespace="c") - parent.add_middleware(ResourceTracingMiddleware("parent", calls)) - - async with Client(parent, mode="legacy") as client: - task = await client.read_resource("item://c/gc/42", task=True) - result = await task.result() - assert result[0].text == "item-42" - - # Middleware runs during task submission, function in Docket - assert calls == [ - "parent:before", - "child:before", - "grandchild:before", - "grandchild:after", - "child:after", - "parent:after", - "grandchild:template", - ] + assert calls == ["parent:before", "parent:after", "grandchild:tool"] -class TestMountedTasksWithTaskMetaParameter: - """Test mounted components called directly with task_meta parameter. +class TestMountedDocketOwnership: + async def test_mounted_child_does_not_own_docket(self, parent_server, child_server): + """The parent owns the Docket; the mounted child does not.""" + async with running_task_server(parent_server): + assert parent_server.docket is not None + assert child_server.docket is None - These tests verify the programmatic API where server.call_tool() or - server.read_resource() is called with an explicit task_meta parameter, - as opposed to using the Client with task=True. - Direct server calls require a running server context, so we use an outer - tool that makes the direct call internally. - """ - - async def test_mounted_tool_with_task_meta_creates_task(self): - """Mounted tool called with task_meta returns CreateTaskResult.""" - from fastmcp.server.tasks.config import TaskMeta - - child = FastMCP("Child") +class TestSlowMountedTaskCancellation: + async def test_cancel_mounted_task(self): + child = FastMCP("child") + release = asyncio.Event() @child.tool(task=True) - async def add(a: int, b: int) -> int: - return a + b + async def slow() -> str: + await release.wait() + return "done" - parent = FastMCP("Parent") + parent = FastMCP("parent") + parent.add_extension(TasksExtension()) parent.mount(child, namespace="child") - @parent.tool - async def outer() -> str: - # Direct call with task_meta from within server context - result = await parent.call_tool( - "child_add", {"a": 2, "b": 3}, task_meta=TaskMeta(ttl=300) + from tests.tasks.task_helpers import cancel_task + + async with running_task_server(parent): + created = await submit_task(parent, "child_slow", {}) + await cancel_task(parent, created.task_id) + release.set() + final = await wait_for_task( + parent, + created.task_id, + target_states=frozenset({"cancelled", "completed"}), ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) - - async def test_mounted_resource_with_task_meta_creates_task(self): - """Mounted resource called with task_meta returns CreateTaskResult.""" - from fastmcp.server.tasks.config import TaskMeta - - child = FastMCP("Child") - - @child.resource("data://info", task=True) - async def get_info() -> str: - return "child info" - - parent = FastMCP("Parent") - parent.mount(child, namespace="child") - - @parent.tool - async def outer() -> str: - result = await parent.read_resource( - "data://child/info", task_meta=TaskMeta(ttl=300) - ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) - - async def test_mounted_template_with_task_meta_creates_task(self): - """Mounted resource template with task_meta returns CreateTaskResult.""" - from fastmcp.server.tasks.config import TaskMeta - - child = FastMCP("Child") - - @child.resource("item://{id}", task=True) - async def get_item(id: str) -> str: - return f"item-{id}" - - parent = FastMCP("Parent") - parent.mount(child, namespace="child") - - @parent.tool - async def outer() -> str: - result = await parent.read_resource( - "item://child/42", task_meta=TaskMeta(ttl=300) - ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) - - async def test_deeply_nested_tool_with_task_meta(self): - """Three-level nested tool works with task_meta.""" - from fastmcp.server.tasks.config import TaskMeta - - grandchild = FastMCP("Grandchild") - - @grandchild.tool(task=True) - async def compute(n: int) -> int: - return n * 3 - - child = FastMCP("Child") - child.mount(grandchild, namespace="gc") - - parent = FastMCP("Parent") - parent.mount(child, namespace="c") - - @parent.tool - async def outer() -> str: - result = await parent.call_tool( - "c_gc_compute", {"n": 7}, task_meta=TaskMeta(ttl=300) - ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) - - async def test_deeply_nested_template_with_task_meta(self): - """Three-level nested template works with task_meta.""" - from fastmcp.server.tasks.config import TaskMeta - - grandchild = FastMCP("Grandchild") - - @grandchild.resource("doc://{name}", task=True) - async def get_doc(name: str) -> str: - return f"doc: {name}" - - child = FastMCP("Child") - child.mount(grandchild, namespace="gc") - - parent = FastMCP("Parent") - parent.mount(child, namespace="c") - - @parent.tool - async def outer() -> str: - result = await parent.read_resource( - "doc://c/gc/readme", task_meta=TaskMeta(ttl=300) - ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) - - async def test_mounted_prompt_with_task_meta_creates_task(self): - """Mounted prompt called with task_meta returns CreateTaskResult.""" - from fastmcp.server.tasks.config import TaskMeta - - child = FastMCP("Child") - - @child.prompt(task=True) - async def greet(name: str) -> str: - return f"Hello, {name}!" - - parent = FastMCP("Parent") - parent.mount(child, namespace="child") - - @parent.tool - async def outer() -> str: - result = await parent.render_prompt( - "child_greet", {"name": "World"}, task_meta=TaskMeta(ttl=300) - ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) - - async def test_deeply_nested_prompt_with_task_meta(self): - """Three-level nested prompt works with task_meta.""" - from fastmcp.server.tasks.config import TaskMeta - - grandchild = FastMCP("Grandchild") - - @grandchild.prompt(task=True) - async def describe(topic: str) -> str: - return f"Information about {topic}" - - child = FastMCP("Child") - child.mount(grandchild, namespace="gc") - - parent = FastMCP("Parent") - parent.mount(child, namespace="c") - - @parent.tool - async def outer() -> str: - result = await parent.render_prompt( - "c_gc_describe", {"topic": "FastMCP"}, task_meta=TaskMeta(ttl=300) - ) - return f"task:{result.task.task_id}" - - async with Client(parent, mode="legacy") as client: - result = await client.call_tool("outer", {}) - assert "task:" in str(result) + assert final.status in {"cancelled", "completed"} diff --git a/tests/tasks/server/test_task_protocol.py b/tests/tasks/server/test_task_protocol.py index 08461bb9a..784e0a51f 100644 --- a/tests/tasks/server/test_task_protocol.py +++ b/tests/tasks/server/test_task_protocol.py @@ -1,85 +1,53 @@ -""" -Tests for SEP-1686 protocol-level task handling. +"""Protocol-level task behavior for SEP-2663 tasks. -Generic protocol tests that use tools as test fixtures. -Tests metadata, notifications, and error handling at the protocol level. +Generic protocol behaviors driven in-process via the task helpers: a submitted +task carries a server-generated id and a TTL, and a task whose tool raises +surfaces its error rather than a result. """ -import pytest +from __future__ import annotations from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp.exceptions import ToolError +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + run_task, + running_task_server, + submit_task, ) -@pytest.fixture -async def task_enabled_server(): - """Create a FastMCP server with task-enabled tools.""" +def _task_server() -> FastMCP: mcp = FastMCP("task-test-server") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def simple_tool(message: str) -> str: - """A simple tool for testing.""" return f"Processed: {message}" @mcp.tool(task=True) async def failing_tool() -> str: - """A tool that always fails.""" - raise ValueError("This tool always fails") + raise ToolError("This tool always fails") return mcp -async def test_task_metadata_includes_task_id_and_ttl(task_enabled_server): - """Task metadata properly includes server-generated taskId and ttl.""" - async with Client(task_enabled_server, mode="legacy") as client: - # Submit with specific ttl (server generates task ID) - task = await client.call_tool( - "simple_tool", - {"message": "test"}, - task=True, - ttl=30000, - ) - assert task - assert not task.returned_immediately - - # Server should have generated a task ID - assert task.task_id is not None - assert isinstance(task.task_id, str) +async def test_task_metadata_includes_task_id_and_ttl(): + """A submitted task carries a server-generated id and a positive TTL.""" + mcp = _task_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "simple_tool", {"message": "test"}) + assert isinstance(created.task_id, str) + assert created.task_id + assert created.ttl_ms is not None and created.ttl_ms > 0 -async def test_task_notification_sent_after_submission(task_enabled_server): - """Server sends an initial task status notification after submission.""" - - @task_enabled_server.tool(task=True) - async def background_tool(message: str) -> str: - return f"Processed: {message}" - - async with Client(task_enabled_server, mode="legacy") as client: - task = await client.call_tool("background_tool", {"message": "test"}, task=True) - assert task - assert not task.returned_immediately - - # Verify we can query the task - status = await task.status() - assert status.task_id == task.task_id - - -async def test_failed_task_stores_error(task_enabled_server): - """Failed tasks store the error in results.""" - - @task_enabled_server.tool(task=True) - async def failing_task_tool() -> str: - raise ValueError("This tool always fails") - - async with Client(task_enabled_server, mode="legacy") as client: - task = await client.call_tool("failing_task_tool", task=True) - assert task - assert not task.returned_immediately - - # Wait for task to fail - status = await task.wait(state="failed", timeout=2.0) - assert status.status == "failed" +async def test_failed_task_stores_error(): + """A task whose tool raises reaches `failed` and stores the error.""" + mcp = _task_server() + async with running_task_server(mcp): + final = await run_task(mcp, "failing_tool", {}) + assert final.status == "failed" + assert final.error is not None + assert "This tool always fails" in final.error["message"] + assert final.result is None diff --git a/tests/tasks/server/test_task_proxy.py b/tests/tasks/server/test_task_proxy.py index 3d4219bac..078e48894 100644 --- a/tests/tasks/server/test_task_proxy.py +++ b/tests/tasks/server/test_task_proxy.py @@ -1,37 +1,54 @@ """ -Tests for MCP SEP-1686 task protocol behavior through proxy servers. +Tests for SEP-2663 task behavior through proxy servers. -Proxy servers explicitly forbid task-augmented execution. All proxy components -(tools, prompts, resources) have task_config.mode="forbidden". - -Clients connecting through proxies can: -- Execute tools/prompts/resources normally (sync execution) -- NOT use task-augmented execution (task=True fails gracefully for tools, - raises MCPError for prompts/resources) +SEP-2663 tasks are tools-only. Proxy servers force every proxied tool to +`task_config.mode="forbidden"`, so a tool that is `task=True` on the backend +runs *synchronously* through the proxy and is never tasked — even when the +client opts the tasks extension in for the request. """ import pytest -from mcp.shared.exceptions import MCPError -from mcp_types import TextContent, TextResourceContents +from docket import Docket +from fastmcp_tasks.models import CreateTaskResult from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.client.transports import FastMCPTransport from fastmcp.server import create_proxy - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp.tools.base import ToolResult +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + _opted_in_request, + auth_scope, + running_task_server, ) +@pytest.fixture(autouse=True) +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to each test's event loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + +async def call_tool_with_optin(server: FastMCP, name: str, arguments: dict): + """Run a `tools/call` with the tasks opt-in bound into the request context.""" + with auth_scope(None), _opted_in_request(name, arguments, None): + return await server.call_tool(name, arguments) + + @pytest.fixture def backend_server() -> FastMCP: - """Create a backend server with task-enabled components. + """A backend server with a task-enabled tool. - The backend has tasks enabled, but the proxy should NOT forward - task execution - it should treat all components as forbidden. + The backend has tasks enabled, but the proxy must NOT forward task + execution — it treats every proxied tool as forbidden. """ mcp = FastMCP("backend-server") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def add_numbers(a: int, b: int) -> int: @@ -43,154 +60,57 @@ def backend_server() -> FastMCP: """Tool that only supports synchronous execution.""" return f"sync: {message}" - @mcp.prompt(task=True) - async def greeting_prompt(name: str) -> str: - """A prompt that can execute as a task.""" - return f"Hello, {name}! Welcome to the system." - - @mcp.resource("data://info.txt", task=True) - async def info_resource() -> str: - """A resource that can be read as a task.""" - return "Important information from the backend" - - @mcp.resource("data://user/{user_id}.json", task=True) - async def user_resource(user_id: str) -> str: - """A resource template that can execute as a task.""" - return f'{{"id": "{user_id}", "name": "User {user_id}"}}' - return mcp @pytest.fixture def proxy_server(backend_server: FastMCP) -> FastMCP: - """Create a proxy server that forwards to the backend.""" - return create_proxy(FastMCPTransport(backend_server)) + """A proxy server that forwards to the backend, with tasks advertised.""" + proxy = create_proxy(FastMCPTransport(backend_server)) + proxy.add_extension(TasksExtension()) + return proxy class TestProxyToolsSyncExecution: - """Test that tools work normally through proxy (sync execution).""" + """Tools work normally through the proxy (synchronous execution).""" async def test_tool_sync_execution_works(self, proxy_server: FastMCP): - """Tool called without task=True works through proxy.""" - async with Client(proxy_server, mode="legacy") as client: + """A tool called without opting in works through the proxy.""" + async with Client(proxy_server) as client: result = await client.call_tool("add_numbers", {"a": 5, "b": 3}) assert "8" in str(result) async def test_sync_only_tool_works(self, proxy_server: FastMCP): - """Sync-only tool works through proxy.""" - async with Client(proxy_server, mode="legacy") as client: + """A sync-only tool works through the proxy.""" + async with Client(proxy_server) as client: result = await client.call_tool("sync_only_tool", {"message": "test"}) assert "sync: test" in str(result) class TestProxyToolsTaskForbidden: - """Test that tools with task=True are forbidden through proxy.""" + """A proxied tool never tasks, even when the client opts in.""" - async def test_tool_task_returns_error_immediately(self, proxy_server: FastMCP): - """Tool called with task=True through proxy returns error immediately.""" - async with Client(proxy_server, mode="legacy") as client: - task = await client.call_tool( - "add_numbers", {"a": 5, "b": 3}, task=True, raise_on_error=False - ) - - # Should return immediately (forbidden behavior) - assert task.returned_immediately - - # Result should be an error - result = await task.result() - assert result.is_error - - async def test_sync_only_tool_task_returns_error_immediately( + async def test_task_enabled_tool_runs_sync_through_proxy( self, proxy_server: FastMCP ): - """Sync-only tool with task=True also returns error immediately.""" - async with Client(proxy_server, mode="legacy") as client: - task = await client.call_tool( - "sync_only_tool", - {"message": "test"}, - task=True, - raise_on_error=False, + """A backend `task=True` tool runs sync through the forbidden proxy.""" + async with running_task_server(proxy_server): + result = await call_tool_with_optin( + proxy_server, "add_numbers", {"a": 5, "b": 3} ) - assert task.returned_immediately - result = await task.result() - assert result.is_error + # The forbidden proxy tool declines to task even with the opt-in. + assert not isinstance(result, CreateTaskResult) + assert isinstance(result, ToolResult) + assert result.structured_content == {"result": 8} + async def test_sync_only_tool_runs_sync_through_proxy(self, proxy_server: FastMCP): + """A sync-only tool also runs sync through the proxy with the opt-in.""" + async with running_task_server(proxy_server): + result = await call_tool_with_optin( + proxy_server, "sync_only_tool", {"message": "test"} + ) -class TestProxyPromptsSyncExecution: - """Test that prompts work normally through proxy (sync execution).""" - - async def test_prompt_sync_execution_works(self, proxy_server: FastMCP): - """Prompt called without task=True works through proxy.""" - async with Client(proxy_server, mode="legacy") as client: - result = await client.get_prompt("greeting_prompt", {"name": "Alice"}) - assert isinstance(result.messages[0].content, TextContent) - assert "Hello, Alice!" in result.messages[0].content.text - - -class TestProxyPromptsTaskForbidden: - """Test that prompts with task=True are forbidden through proxy.""" - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_prompt_task_raises_mcp_error(self, proxy_server: FastMCP): - """Prompt called with task=True through proxy raises MCPError.""" - async with Client(proxy_server, mode="legacy") as client: - with pytest.raises(MCPError) as exc_info: - await client.get_prompt("greeting_prompt", {"name": "Alice"}, task=True) - - assert "does not support task-augmented execution" in str(exc_info.value) - - -class TestProxyResourcesSyncExecution: - """Test that resources work normally through proxy (sync execution).""" - - async def test_resource_sync_execution_works(self, proxy_server: FastMCP): - """Resource read without task=True works through proxy.""" - async with Client(proxy_server, mode="legacy") as client: - result = await client.read_resource("data://info.txt") - assert isinstance(result[0], TextResourceContents) - assert "Important information from the backend" in result[0].text - - async def test_resource_template_sync_execution_works(self, proxy_server: FastMCP): - """Resource template without task=True works through proxy.""" - async with Client(proxy_server, mode="legacy") as client: - result = await client.read_resource("data://user/42.json") - assert isinstance(result[0], TextResourceContents) - assert '"id": "42"' in result[0].text - - -class TestProxyResourcesTaskForbidden: - """Test that resources with task=True are forbidden through proxy.""" - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_resource_task_raises_mcp_error(self, proxy_server: FastMCP): - """Resource read with task=True through proxy raises MCPError.""" - async with Client(proxy_server, mode="legacy") as client: - with pytest.raises(MCPError) as exc_info: - await client.read_resource("data://info.txt", task=True) - - assert "does not support task-augmented execution" in str(exc_info.value) - - @pytest.mark.xfail( - reason="SDK v2 has no `task` field on GetPromptRequestParams / " - "ReadResourceRequestParams; prompt/resource task submission is not " - "wire-expressible and always graceful-degrades (sdk-feedback #3).", - strict=True, - ) - async def test_resource_template_task_raises_mcp_error(self, proxy_server: FastMCP): - """Resource template with task=True through proxy raises MCPError.""" - async with Client(proxy_server, mode="legacy") as client: - with pytest.raises(MCPError) as exc_info: - await client.read_resource("data://user/42.json", task=True) - - assert "does not support task-augmented execution" in str(exc_info.value) + assert not isinstance(result, CreateTaskResult) + assert isinstance(result, ToolResult) + assert result.structured_content == {"result": "sync: test"} diff --git a/tests/tasks/server/test_task_return_types.py b/tests/tasks/server/test_task_return_types.py index ddd80777c..e1a3106c6 100644 --- a/tests/tasks/server/test_task_return_types.py +++ b/tests/tasks/server/test_task_return_types.py @@ -1,9 +1,11 @@ """ -Tests to verify all return types work identically with task=True. +Tests to verify all tool return types work identically with task=True. -These tests ensure that enabling background task support doesn't break -existing functionality - any tool/prompt/resource should work exactly -the same whether task=True or task=False. +SEP-2663 tasks are tools-only. Every tool below is exercised twice: once +synchronously (no tasks opt-in) and once as a background task. Both paths run +the same `tool.convert_result(...).to_mcp_result()` pipeline, so the inlined +task result must be byte-for-byte identical to the synchronous result. These +tests assert that equivalence across every supported return type. """ from dataclasses import dataclass @@ -12,19 +14,66 @@ from pathlib import Path from typing import Any from uuid import UUID +import mcp_types import pytest +from docket import Docket from pydantic import BaseModel from typing_extensions import TypedDict from fastmcp import FastMCP -from fastmcp.client import Client +from fastmcp.tools.base import ToolResult from fastmcp.utilities.types import Audio, File, Image - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + call_tool_without_optin, + run_task, + running_task_server, ) +@pytest.fixture(autouse=True) +def reset_docket_memory_server(): + """Force a fresh memory:// Docket server bound to each test's event loop.""" + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + yield + if hasattr(Docket, "_memory_server"): + delattr(Docket, "_memory_server") + + +def _sync_result_to_wire(result: ToolResult) -> dict[str, Any]: + """Serialize a synchronous ToolResult into the inlined task wire shape.""" + mcp_result = result.to_mcp_result() + if isinstance(mcp_result, mcp_types.CallToolResult): + call_tool_result = mcp_result + elif isinstance(mcp_result, tuple): + content, structured_content = mcp_result + call_tool_result = mcp_types.CallToolResult( + content=content, + structuredContent=structured_content, + ) + else: + call_tool_result = mcp_types.CallToolResult(content=mcp_result) + return call_tool_result.model_dump(by_alias=True, mode="json", exclude_none=True) + + +async def assert_task_matches_sync( + server: FastMCP, + tool_name: str, + arguments: dict[str, Any] | None = None, +) -> None: + """Run a tool sync and as a task; assert the inlined results are identical.""" + async with running_task_server(server): + sync_result = await call_tool_without_optin(server, tool_name, arguments) + assert isinstance(sync_result, ToolResult) + + task_result = await run_task(server, tool_name, arguments) + assert task_result.status == "completed" + assert task_result.result is not None + + assert task_result.result == _sync_result_to_wire(sync_result) + + class UserData(BaseModel): """Example structured output.""" @@ -33,47 +82,45 @@ class UserData(BaseModel): active: bool -@pytest.fixture -async def return_type_server(): - """Server with tools that return various types.""" - mcp = FastMCP("return-type-test") +# ============================================================================== +# Basic Types +# ============================================================================== + + +@pytest.fixture +def return_type_server(): + """Server with tools that return various basic types.""" + mcp = FastMCP("return-type-test") + mcp.add_extension(TasksExtension()) - # String return @mcp.tool(task=True) async def return_string() -> str: return "Hello, World!" - # Integer return @mcp.tool(task=True) async def return_int() -> int: return 42 - # Float return @mcp.tool(task=True) async def return_float() -> float: return 3.14159 - # Boolean return @mcp.tool(task=True) async def return_bool() -> bool: return True - # Dict return @mcp.tool(task=True) async def return_dict() -> dict[str, int]: return {"count": 100, "total": 500} - # List return @mcp.tool(task=True) async def return_list() -> list[str]: return ["apple", "banana", "cherry"] - # BaseModel return (structured output) @mcp.tool(task=True) async def return_model() -> UserData: return UserData(name="Alice", age=30, active=True) - # None/null return @mcp.tool(task=True) async def return_none() -> None: return None @@ -82,150 +129,24 @@ async def return_type_server(): @pytest.mark.parametrize( - "tool_name,expected_type,expected_value", + "tool_name", [ - ("return_string", str, "Hello, World!"), - ("return_int", int, 42), - ("return_float", float, 3.14159), - ("return_bool", bool, True), - ("return_dict", dict, {"count": 100, "total": 500}), - ("return_list", list, ["apple", "banana", "cherry"]), - ("return_none", type(None), None), + "return_string", + "return_int", + "return_float", + "return_bool", + "return_dict", + "return_list", + "return_model", + "return_none", ], ) -async def test_task_basic_types( +async def test_task_basic_types_match_sync( return_type_server: FastMCP, tool_name: str, - expected_type: type, - expected_value: Any, ): - """Task mode returns basic types correctly.""" - async with Client(return_type_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert isinstance(result.data, expected_type) - assert result.data == expected_value - - -async def test_task_model_return(return_type_server): - """Task mode returns same BaseModel (as dict) as immediate mode.""" - async with Client(return_type_server, mode="legacy") as client: - task = await client.call_tool("return_model", task=True) - result = await task - - # Client deserializes to dynamic class (type name lost with title pruning) - assert result.data.__class__.__name__ == "Root" - assert result.data.name == "Alice" - assert result.data.age == 30 - assert result.data.active is True - - -async def test_task_vs_immediate_equivalence(return_type_server): - """Verify task mode and immediate mode return identical results.""" - async with Client(return_type_server, mode="legacy") as client: - # Test a few types to verify equivalence - tools_to_test = ["return_string", "return_int", "return_dict"] - - for tool_name in tools_to_test: - # Call as task - task = await client.call_tool(tool_name, task=True) - task_result = await task - - # Call immediately (server should decline background execution when no task meta) - immediate_result = await client.call_tool(tool_name) - - # Results should be identical - assert task_result.data == immediate_result.data, ( - f"Mismatch for {tool_name}" - ) - - -@pytest.fixture -async def prompt_return_server(): - """Server with prompts that return various message structures.""" - mcp = FastMCP("prompt-return-test") - - @mcp.prompt(task=True) - async def single_message_prompt() -> str: - """Return a single string message.""" - return "Single message content" - - @mcp.prompt(task=True) - async def multi_message_prompt() -> list[str]: - """Return multiple messages.""" - return [ - "First message", - "Second message", - "Third message", - ] - - return mcp - - -async def test_prompt_task_single_message(prompt_return_server): - """Prompt task returns single message correctly.""" - async with Client(prompt_return_server, mode="legacy") as client: - task = await client.get_prompt("single_message_prompt", task=True) - result = await task - - assert len(result.messages) == 1 - assert result.messages[0].content.text == "Single message content" - - -async def test_prompt_task_multiple_messages(prompt_return_server): - """Prompt task returns multiple messages correctly.""" - async with Client(prompt_return_server, mode="legacy") as client: - task = await client.get_prompt("multi_message_prompt", task=True) - result = await task - - assert len(result.messages) == 3 - assert result.messages[0].content.text == "First message" - assert result.messages[1].content.text == "Second message" - assert result.messages[2].content.text == "Third message" - - -@pytest.fixture -async def resource_return_server(): - """Server with resources that return various content types.""" - mcp = FastMCP("resource-return-test") - - @mcp.resource("text://simple", task=True) - async def simple_text() -> str: - """Return simple text content.""" - return "Simple text resource" - - @mcp.resource("data://json", task=True) - async def json_data() -> str: - """Return JSON-like data.""" - import json - - return json.dumps({"key": "value", "count": 123}) - - return mcp - - -async def test_resource_task_text_content(resource_return_server): - """Resource task returns text content correctly.""" - async with Client(resource_return_server, mode="legacy") as client: - task = await client.read_resource("text://simple", task=True) - contents = await task - - assert len(contents) == 1 - assert contents[0].text == "Simple text resource" - - -async def test_resource_task_json_content(resource_return_server): - """Resource task returns structured content correctly.""" - async with Client(resource_return_server, mode="legacy") as client: - task = await client.read_resource("data://json", task=True) - contents = await task - - # Content should be JSON serialized - assert len(contents) == 1 - import json - - data = json.loads(contents[0].text) - assert data == {"key": "value", "count": 123} + """Task mode returns basic types identically to the synchronous path.""" + await assert_task_matches_sync(return_type_server, tool_name) # ============================================================================== @@ -234,9 +155,10 @@ async def test_resource_task_json_content(resource_return_server): @pytest.fixture -async def binary_type_server(): +def binary_type_server(tmp_path): """Server with tools returning binary and special types.""" mcp = FastMCP("binary-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def return_bytes() -> bytes: @@ -258,44 +180,15 @@ async def binary_type_server(): @pytest.mark.parametrize( - "tool_name,expected_type,assertion_fn", - [ - ( - "return_bytes", - type(None), - lambda r: ( - r.data is None and any("Hello bytes!" in c.text for c in r.content) - ), - ), - ( - "return_uuid", - str, - lambda r: r.data == "12345678-1234-5678-1234-567812345678", - ), - ( - "return_path", - str, - lambda r: "tmp" in r.data and "test.txt" in r.data, - ), - ( - "return_datetime", - datetime, - lambda r: r.data == datetime(2025, 11, 5, 12, 30, 45), - ), - ], + "tool_name", + ["return_bytes", "return_uuid", "return_path", "return_datetime"], ) -async def test_task_binary_types( +async def test_task_binary_types_match_sync( binary_type_server: FastMCP, tool_name: str, - expected_type: type, - assertion_fn: Any, ): - """Task mode handles binary and special types.""" - async with Client(binary_type_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert isinstance(result.data, expected_type) - assert assertion_fn(result) + """Task mode handles binary and special types identically to sync.""" + await assert_task_matches_sync(binary_type_server, tool_name) # ============================================================================== @@ -304,9 +197,10 @@ async def test_task_binary_types( @pytest.fixture -async def collection_server(): +def collection_server(): """Server with tools returning various collection types.""" mcp = FastMCP("collection-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def return_tuple() -> tuple[int, str, bool]: @@ -328,36 +222,15 @@ async def collection_server(): @pytest.mark.parametrize( - "tool_name,expected_type,expected_value", - [ - ("return_tuple", list, [42, "hello", True]), - ("return_set", set, {1, 2, 3}), - ("return_empty_list", list, []), - ], + "tool_name", + ["return_tuple", "return_set", "return_empty_list", "return_empty_dict"], ) -async def test_task_collection_types( +async def test_task_collection_types_match_sync( collection_server: FastMCP, tool_name: str, - expected_type: type, - expected_value: Any, ): - """Task mode handles collection types.""" - async with Client(collection_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert isinstance(result.data, expected_type) - assert result.data == expected_value - - -async def test_task_empty_dict_return(collection_server): - """Task mode handles empty dict return.""" - async with Client(collection_server, mode="legacy") as client: - task = await client.call_tool("return_empty_dict", task=True) - result = await task - # Empty structured content becomes None in data - assert result.data is None - # But structured content is still {} - assert result.structured_content == {} + """Task mode handles collection types identically to sync.""" + await assert_task_matches_sync(collection_server, tool_name) # ============================================================================== @@ -366,11 +239,11 @@ async def test_task_empty_dict_return(collection_server): @pytest.fixture -async def media_server(tmp_path): +def media_server(tmp_path): """Server with tools returning media types.""" mcp = FastMCP("media-test") + mcp.add_extension(TasksExtension()) - # Create test files test_image = tmp_path / "test.png" test_image.write_bytes(b"\x89PNG\r\n\x1a\n" + b"fake png data") @@ -400,40 +273,15 @@ async def media_server(tmp_path): @pytest.mark.parametrize( - "tool_name,assertion_fn", - [ - ( - "return_image_path", - lambda r: len(r.content) == 1 and r.content[0].type == "image", - ), - ( - "return_image_data", - lambda r: ( - len(r.content) == 1 - and r.content[0].type == "image" - and r.content[0].mime_type == "image/png" - ), - ), - ( - "return_audio", - lambda r: len(r.content) == 1 and r.content[0].type in ["text", "audio"], - ), - ( - "return_file", - lambda r: len(r.content) == 1 and r.content[0].type == "resource", - ), - ], + "tool_name", + ["return_image_path", "return_image_data", "return_audio", "return_file"], ) -async def test_task_media_types( +async def test_task_media_types_match_sync( media_server: FastMCP, tool_name: str, - assertion_fn: Any, ): - """Task mode handles media types (Image, Audio, File).""" - async with Client(media_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert assertion_fn(result) + """Task mode handles media types (Image, Audio, File) identically to sync.""" + await assert_task_matches_sync(media_server, tool_name) # ============================================================================== @@ -457,9 +305,10 @@ class PersonDataclass: @pytest.fixture -async def structured_type_server(): +def structured_type_server(): """Server with tools returning structured types.""" mcp = FastMCP("structured-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def return_typeddict() -> PersonTypedDict: @@ -489,67 +338,22 @@ async def structured_type_server(): @pytest.mark.parametrize( - "tool_name,expected_name,expected_age", + "tool_name", [ - ("return_typeddict", "Bob", 25), - ("return_dataclass", "Charlie", 35), + "return_typeddict", + "return_dataclass", + "return_union", + "return_union_int", + "return_optional", + "return_optional_none", ], ) -async def test_task_structured_dict_types( +async def test_task_structured_types_match_sync( structured_type_server: FastMCP, tool_name: str, - expected_name: str, - expected_age: int, ): - """Task mode handles TypedDict and dataclass returns.""" - async with Client(structured_type_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - # Both deserialize to dynamic Root class - assert result.data.name == expected_name - assert result.data.age == expected_age - - -@pytest.mark.parametrize( - "tool_name,expected_type,expected_value", - [ - ("return_union", str, "string value"), - ("return_union_int", int, 123), - ], -) -async def test_task_union_types( - structured_type_server: FastMCP, - tool_name: str, - expected_type: type, - expected_value: Any, -): - """Task mode handles union type branches.""" - async with Client(structured_type_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert isinstance(result.data, expected_type) - assert result.data == expected_value - - -@pytest.mark.parametrize( - "tool_name,expected_type,expected_value", - [ - ("return_optional", str, "has value"), - ("return_optional_none", type(None), None), - ], -) -async def test_task_optional_types( - structured_type_server: FastMCP, - tool_name: str, - expected_type: type, - expected_value: Any, -): - """Task mode handles Optional types.""" - async with Client(structured_type_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert isinstance(result.data, expected_type) - assert result.data == expected_value + """Task mode handles TypedDict, dataclass, union and optional returns.""" + await assert_task_matches_sync(structured_type_server, tool_name) # ============================================================================== @@ -558,7 +362,7 @@ async def test_task_optional_types( @pytest.fixture -async def mcp_content_server(tmp_path): +def mcp_content_server(tmp_path): """Server with tools returning MCP content blocks.""" import base64 @@ -571,6 +375,7 @@ async def mcp_content_server(tmp_path): ) mcp = FastMCP("content-test") + mcp.add_extension(TasksExtension()) test_image = tmp_path / "content.png" test_image.write_bytes(b"\x89PNG\r\n\x1a\n" + b"content") @@ -616,58 +421,18 @@ async def mcp_content_server(tmp_path): @pytest.mark.parametrize( - "tool_name,assertion_fn", + "tool_name", [ - ( - "return_text_content", - lambda r: ( - len(r.content) == 1 - and r.content[0].type == "text" - and r.content[0].text == "Direct text content" - ), - ), - ( - "return_image_content", - lambda r: ( - len(r.content) == 1 - and r.content[0].type == "image" - and r.content[0].mime_type == "image/png" - ), - ), - ( - "return_embedded_resource", - lambda r: len(r.content) == 1 and r.content[0].type == "resource", - ), - ( - "return_resource_link", - lambda r: ( - len(r.content) == 1 - and r.content[0].type == "resource_link" - and str(r.content[0].uri) == "test://linked" - ), - ), + "return_text_content", + "return_image_content", + "return_embedded_resource", + "return_resource_link", + "return_mixed_content", ], ) -async def test_task_mcp_content_types( +async def test_task_mcp_content_types_match_sync( mcp_content_server: FastMCP, tool_name: str, - assertion_fn: Any, ): - """Task mode handles MCP content block types.""" - async with Client(mcp_content_server, mode="legacy") as client: - task = await client.call_tool(tool_name, task=True) - result = await task - assert assertion_fn(result) - - -async def test_task_mixed_content_return(mcp_content_server): - """Task mode handles mixed content list return.""" - async with Client(mcp_content_server, mode="legacy") as client: - task = await client.call_tool("return_mixed_content", task=True) - result = await task - assert len(result.content) == 3 - assert result.content[0].type == "text" - assert result.content[0].text == "First block" - assert result.content[1].type == "image" - assert result.content[2].type == "text" - assert result.content[2].text == "Third block" + """Task mode handles MCP content block types identically to sync.""" + await assert_task_matches_sync(mcp_content_server, tool_name) diff --git a/tests/tasks/server/test_task_security.py b/tests/tasks/server/test_task_security.py index 2ee18c6db..088f52153 100644 --- a/tests/tasks/server/test_task_security.py +++ b/tests/tasks/server/test_task_security.py @@ -1,153 +1,104 @@ -""" -Tests for authorization-based task isolation (CRITICAL SECURITY). +"""Authorization-based task isolation (CRITICAL SECURITY). -Ensures that tasks are properly scoped to authorization identity and clients -cannot access each other's tasks. +Tasks are scoped to the caller's authorization identity via the auth-scoped +compound Docket key, so a caller can only resolve tasks it created. A cross-scope +task id is indistinguishable from a missing one (-32602 "not found"), which keeps +task existence from leaking across callers. These tests drive the task lifecycle +in-process (there is no client task API until Phase 4), binding a different +access token per caller through the shared helper. """ +from __future__ import annotations + import pytest -from mcp.server.auth.middleware.auth_context import ( - auth_context_var, -) -from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser +from mcp.shared.exceptions import MCPError from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.server.auth import AccessToken - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + get_task, + make_access_token, + run_task, + running_task_server, + submit_task, + wait_for_task, ) @pytest.fixture -def task_server(): - """Create a server with background tasks enabled.""" +def task_server() -> FastMCP: mcp = FastMCP("security-test-server") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def secret_tool(data: str) -> str: - """A tool that processes sensitive data.""" return f"Secret result: {data}" return mcp async def test_same_client_can_access_all_its_tasks(task_server: FastMCP): - """A single authenticated client can access all tasks it created.""" - token = AccessToken( - token="token-a", - client_id="client-a", - scopes=["read"], - ) - reset = auth_context_var.set(AuthenticatedUser(token)) - try: - async with Client(task_server, mode="legacy") as client: - task1 = await client.call_tool( - "secret_tool", {"data": "first"}, task=True, task_id="task-1" - ) - task2 = await client.call_tool( - "secret_tool", {"data": "second"}, task=True, task_id="task-2" - ) - - await task1.wait(timeout=2.0) - await task2.wait(timeout=2.0) - - result1 = await task1.result() - result2 = await task2.result() - - assert "first" in str(result1.data) - assert "second" in str(result2.data) - finally: - auth_context_var.reset(reset) + """A single authenticated caller can resolve every task it created.""" + token = make_access_token("client-a") + async with running_task_server(task_server): + first = await run_task( + task_server, "secret_tool", {"data": "first"}, access_token=token + ) + second = await run_task( + task_server, "secret_tool", {"data": "second"}, access_token=token + ) + assert "first" in first.result["content"][0]["text"] + assert "second" in second.result["content"][0]["text"] async def test_unauthenticated_client_can_access_its_tasks(task_server: FastMCP): - """An unauthenticated client can access tasks it created (by task ID).""" - async with Client(task_server, mode="legacy") as client: - task = await client.call_tool( - "secret_tool", {"data": "hello"}, task=True, task_id="my-task" - ) - await task.wait(timeout=2.0) - result = await task.result() - assert "hello" in str(result.data) - - -def _set_auth(client_id: str, sub: str | None = None): - """Install an auth context for a given client_id/sub. Returns the reset token.""" - claims = {"sub": sub} if sub else {} - token = AccessToken( - token=f"token-{client_id}-{sub or ''}", - client_id=client_id, - scopes=["read"], - claims=claims, - ) - return auth_context_var.set(AuthenticatedUser(token)) - - -async def _submit_task_id(client: Client, data: str) -> str: - """Submit a background task and return its server-assigned task id.""" - task = await client.call_tool("secret_tool", {"data": data}, task=True) - await task.wait(timeout=2.0) - return task.task_id + """An anonymous caller can resolve tasks in the anonymous keyspace.""" + async with running_task_server(task_server): + final = await run_task(task_server, "secret_tool", {"data": "hello"}) + assert "hello" in final.result["content"][0]["text"] async def test_distinct_clients_cannot_access_each_others_tasks( task_server: FastMCP, ): - """Two distinct authenticated clients live in disjoint scopes — looking up - a peer's task id returns 'not found'.""" - reset = _set_auth("client-a") - try: - async with Client(task_server, mode="legacy") as client_a: - task_id = await _submit_task_id(client_a, "client-a-secret") - finally: - auth_context_var.reset(reset) - - reset = _set_auth("client-b") - try: - async with Client(task_server, mode="legacy") as client_b: - with pytest.raises(Exception, match="not found"): - await client_b.get_task_status(task_id) - finally: - auth_context_var.reset(reset) + """Two distinct client_ids live in disjoint scopes: a peer's id is 'not found'.""" + alice = make_access_token("client-a") + bob = make_access_token("client-b") + async with running_task_server(task_server): + created = await submit_task( + task_server, "secret_tool", {"data": "a-secret"}, access_token=alice + ) + with pytest.raises(MCPError, match="not found"): + await get_task(task_server, created.task_id, access_token=bob) async def test_distinct_subs_same_client_id_cannot_access_each_others_tasks( task_server: FastMCP, ): - """Fixed-OAuth case: two users share a client_id but have distinct ``sub`` - claims. The ``sub``-aware scope must still isolate them.""" - shared_client = "shared-oauth-app" - - reset = _set_auth(shared_client, sub="user-alice") - try: - async with Client(task_server, mode="legacy") as alice: - task_id = await _submit_task_id(alice, "alice-secret") - finally: - auth_context_var.reset(reset) - - reset = _set_auth(shared_client, sub="user-bob") - try: - async with Client(task_server, mode="legacy") as bob: - with pytest.raises(Exception, match="not found"): - await bob.get_task_status(task_id) - finally: - auth_context_var.reset(reset) + """Fixed-OAuth case: one client_id, distinct ``sub`` claims stay isolated.""" + shared = "shared-oauth-app" + alice = make_access_token(shared, sub="user-alice") + bob = make_access_token(shared, sub="user-bob") + async with running_task_server(task_server): + created = await submit_task( + task_server, "secret_tool", {"data": "alice-secret"}, access_token=alice + ) + with pytest.raises(MCPError, match="not found"): + await get_task(task_server, created.task_id, access_token=bob) async def test_authenticated_and_anonymous_keyspaces_are_disjoint( task_server: FastMCP, ): - """An anonymous client must not be able to read an authenticated client's - tasks (and vice versa) even when colliding on task id.""" - reset = _set_auth("client-a") - try: - async with Client(task_server, mode="legacy") as authed: - authed_task_id = await _submit_task_id(authed, "authed-secret") - finally: - auth_context_var.reset(reset) - - async with Client(task_server, mode="legacy") as anon: - with pytest.raises(Exception, match="not found"): - await anon.get_task_status(authed_task_id) + """An anonymous caller cannot read an authenticated caller's task.""" + authed = make_access_token("client-a") + async with running_task_server(task_server): + created = await submit_task( + task_server, "secret_tool", {"data": "authed-secret"}, access_token=authed + ) + # No access_token -> anonymous keyspace -> cannot resolve the authed task. + with pytest.raises(MCPError, match="not found"): + await get_task(task_server, created.task_id) + # And the authenticated caller still resolves it. + seen = await wait_for_task(task_server, created.task_id, access_token=authed) + assert seen.status == "completed" diff --git a/tests/tasks/server/test_task_status_notifications.py b/tests/tasks/server/test_task_status_notifications.py deleted file mode 100644 index 3487148f7..000000000 --- a/tests/tasks/server/test_task_status_notifications.py +++ /dev/null @@ -1,168 +0,0 @@ -""" -Tests for notifications/tasks/status subscription mechanism (SEP-1686 lines 436-444). - -Per the spec, servers MAY send notifications/tasks/status when task state changes. -This is an optional optimization that reduces client polling frequency. - -These tests verify that the subscription mechanism works correctly without breaking -existing functionality. Notification delivery is best-effort and clients MUST NOT -rely on receiving them. -""" - -import asyncio - -import pytest - -from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" -) - - -@pytest.fixture -async def notification_server(): - """Create a server for testing task status notifications.""" - mcp = FastMCP("notification-test") - - @mcp.tool(task=True) - async def quick_task(value: int) -> int: - """Quick task that completes immediately.""" - return value * 2 - - @mcp.tool(task=True) - async def slow_task() -> str: - """Task that never completes on its own. - - Only used to verify disconnect-while-running doesn't crash - the - test disconnects before the task would finish, so it never needs - to actually complete. - """ - await asyncio.Event().wait() - return "completed" - - @mcp.tool(task=True) - async def failing_task() -> str: - """Task that always fails.""" - raise ValueError("Task failed intentionally") - - @mcp.prompt(task=True) - async def test_prompt(name: str) -> str: - """Test prompt for background execution.""" - return f"Hello, {name}!" - - @mcp.resource("test://resource", task=True) - async def test_resource() -> str: - """Test resource for background execution.""" - return "resource content" - - return mcp - - -async def test_subscription_spawned_for_tool_task(notification_server: FastMCP): - """Subscription task is spawned when tool task is created.""" - async with Client(notification_server, mode="legacy") as client: - # Create task - should spawn subscription - task = await client.call_tool("quick_task", {"value": 5}, task=True) - - # Task should complete normally - result = await task - assert result.data == 10 - - # Subscription should clean up automatically - # (No way to directly test, but shouldn't cause issues) - - -async def test_subscription_handles_task_completion(notification_server: FastMCP): - """Subscription properly handles task completion and cleanup.""" - async with Client(notification_server, mode="legacy") as client: - # Multiple tasks should each get their own subscription - task1 = await client.call_tool("quick_task", {"value": 1}, task=True) - task2 = await client.call_tool("quick_task", {"value": 2}, task=True) - task3 = await client.call_tool("quick_task", {"value": 3}, task=True) - - # All should complete successfully - result1 = await task1 - result2 = await task2 - result3 = await task3 - - assert result1.data == 2 - assert result2.data == 4 - assert result3.data == 6 - - # Subscriptions clean up deterministically via the connection's - # exit stack when the client disconnects (see test below), so no - # settling wait is needed here. - - -async def test_subscription_handles_task_failure(notification_server: FastMCP): - """Subscription properly handles task failure.""" - async with Client(notification_server, mode="legacy") as client: - task = await client.call_tool("failing_task", {}, task=True) - - # Task should fail - with pytest.raises(Exception): - await task - - # Subscription cleans up deterministically via the connection's - # exit stack on disconnect; no settling wait is needed here. - - -async def test_subscription_for_prompt_tasks(notification_server: FastMCP): - """Subscriptions work for prompt tasks.""" - async with Client(notification_server, mode="legacy") as client: - task = await client.get_prompt("test_prompt", {"name": "World"}, task=True) - - result = await task - # Prompt result has messages - assert result - - # Subscription cleans up deterministically via the connection's - # exit stack on disconnect; no settling wait is needed here. - - -async def test_subscription_for_resource_tasks(notification_server: FastMCP): - """Subscriptions work for resource tasks.""" - async with Client(notification_server, mode="legacy") as client: - task = await client.read_resource("test://resource", task=True) - - result = await task - assert result # Resource contents - - # Subscription cleans up deterministically via the connection's - # exit stack on disconnect; no settling wait is needed here. - - -async def test_subscriptions_cleanup_on_session_disconnect( - notification_server: FastMCP, -): - """Subscriptions are cleaned up when session disconnects.""" - # Start session and create task - # Task submission is a handshake-era capability, so this pins the legacy era. - async with Client(notification_server, mode="legacy") as client: - task = await client.call_tool("slow_task", {}, task=True) - task_id = task.task_id - # Disconnect before task completes (session __aexit__ cancels subscriptions) - - # Session is now closed, subscription should be cancelled - # Task continues in Docket but notification subscription is gone - # This test passing means no crash occurred during cleanup - assert task_id # Task was created - - -async def test_multiple_concurrent_subscriptions(notification_server: FastMCP): - """Multiple concurrent tasks each have their own subscription.""" - async with Client(notification_server, mode="legacy") as client: - # Start many tasks concurrently - tasks = [] - for i in range(10): - task = await client.call_tool("quick_task", {"value": i}, task=True) - tasks.append(task) - - # All should complete - results = await asyncio.gather(*tasks) - assert len(results) == 10 - - # All subscriptions clean up deterministically via the connection's - # exit stack on disconnect; no settling wait is needed here. diff --git a/tests/tasks/server/test_task_tools.py b/tests/tasks/server/test_task_tools.py index 15b4358c9..31aa1666f 100644 --- a/tests/tasks/server/test_task_tools.py +++ b/tests/tasks/server/test_task_tools.py @@ -1,60 +1,77 @@ -""" -Tests for server-side tool task behavior. +"""Server-side tool task behavior for SEP-2663 tasks. -Tests tool-specific task handling, parallel to test_task_prompts.py -and test_task_resources.py. +Covers task=True/False decoration, argument coercion parity between the +synchronous and task-submission paths (including the strict-validation flag), +immediate task metadata on submission, background execution with status polling, +and the rule that a forbidden (task=False) tool runs synchronously even when the +caller opts into tasks. Driven in-process via the task helpers because there is +no client task-submission API until Phase 4. """ +from __future__ import annotations + import asyncio import functools -import mcp_types import pytest -from fastmcp_tasks.client import ToolTask +from fastmcp_tasks.models import CreateTaskResult from pydantic import BaseModel from fastmcp import FastMCP -from fastmcp.client import Client -from fastmcp.client.messages import MessageHandler -from fastmcp.exceptions import ToolError +from fastmcp.exceptions import ValidationError from fastmcp.tools.function_tool import _resolve_param_hints - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + _opted_in_request, + auth_scope, + call_tool_without_optin, + get_task, + run_task, + running_task_server, + submit_task, + wait_for_task, ) -@pytest.fixture -async def tool_server(): - """Create a FastMCP server with task-enabled tools.""" - mcp = FastMCP("tool-task-server") - - @mcp.tool(task=True) - async def simple_tool(message: str) -> str: - """A simple tool for testing.""" - return f"Processed: {message}" - - @mcp.tool(task=False) - async def sync_only_tool(message: str) -> str: - """Tool with task=False.""" - return f"Sync: {message}" - - return mcp - - class _Item(BaseModel): value: str -async def test_task_tool_validates_model_arguments(): - """Model-typed args are coerced to model instances for task calls (#4349). +async def _opted_in_call(server: FastMCP, name: str, arguments: dict | None = None): + """Run a `tools/call` WITH the tasks opt-in bound (used to prove sync paths).""" + with auth_scope(None), _opted_in_request(name, arguments or {}, None): + return await server.call_tool(name, arguments or {}) - The synchronous path validates arguments through the function's - TypeAdapter, so a parameter typed as a Pydantic model arrives as a model - instance. The task path must coerce the same way rather than passing the - raw dict through to the function. + +def _tool_server() -> FastMCP: + mcp = FastMCP("tool-task-server") + mcp.add_extension(TasksExtension()) + + @mcp.tool(task=True) + async def simple_tool(message: str) -> str: + return f"Processed: {message}" + + @mcp.tool(task=False) + async def sync_only_tool(message: str) -> str: + return f"Sync: {message}" + + return mcp + + +# --------------------------------------------------------------------------- +# Argument coercion parity +# --------------------------------------------------------------------------- + + +async def test_task_tool_coerces_model_arguments(): + """Model-typed args are coerced to model instances on the task path (#4349). + + The synchronous path validates arguments through the function's TypeAdapter, + so a parameter typed as a Pydantic model arrives as a model instance. The + task path must coerce identically rather than passing the raw dict through. """ mcp = FastMCP("tool-task-validation-server") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def inspect_items(item: _Item, items: list[_Item]) -> dict[str, str]: @@ -62,102 +79,65 @@ async def test_task_tool_validates_model_arguments(): arguments = {"item": {"value": "a"}, "items": [{"value": "b"}]} expected = {"item": "_Item", "element": "_Item"} + async with running_task_server(mcp): + sync_result = await call_tool_without_optin(mcp, "inspect_items", arguments) + final = await run_task(mcp, "inspect_items", arguments) - async with Client(mcp, mode="legacy") as client: - sync_result = await client.call_tool("inspect_items", arguments) - task = await client.call_tool("inspect_items", arguments, task=True) - task_result = await task.result() - - assert sync_result.data == expected - assert task_result.data == expected + assert sync_result.structured_content == expected + assert final.result["structuredContent"] == expected -async def test_task_tool_invalid_arguments_fail_before_task_state(): - """Invalid task arguments are rejected before any task state is created. +async def test_task_arguments_are_coerced_like_sync_path(): + """A string-for-int arg coerces on the task path exactly as on the sync path.""" + mcp = FastMCP("coerce-task-server") + mcp.add_extension(TasksExtension()) - Coercion runs up front in submit_to_docket, so a validation failure surfaces - before the task's Redis metadata and initial "working" status notification - are written. Otherwise an invalid input would orphan a task the client had - already observed via that notification. - """ + @mcp.tool(task=True) + async def square(n: int) -> int: + return n * n - class _Recorder(MessageHandler): - def __init__(self): - super().__init__() - self.methods: list[str] = [] - - async def on_notification(self, message: mcp_types.ServerNotification) -> None: - self.methods.append(message.method) - - server = FastMCP("tool-task-invalid-args-server") - - @server.tool(task=True) - async def needs_item(item: _Item) -> str: - return item.value - - recorder = _Recorder() - async with Client(server, mode="legacy", message_handler=recorder) as client: - # `item` is missing its required `value` field. - task = await client.call_tool("needs_item", {"item": {}}, task=True) - assert task.returned_immediately - with pytest.raises(ToolError): - await task.result() - - assert "notifications/tasks/status" not in recorder.methods + async with running_task_server(mcp): + final = await run_task(mcp, "square", {"n": "1"}) + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": 1} async def test_task_submission_honors_strict_input_validation(): - """Strict input validation applies to task submissions, not just sync calls. + """Strict input validation rejects lax coercion on the task path too. - With ``strict_input_validation=True``, a lax coercion like ``{"n": "1"}`` - for an ``int`` parameter is rejected on the synchronous path. The task - submission path must reject it identically rather than silently coercing - and queueing it — otherwise ``task=True`` would bypass the strict flag. + With ``strict_input_validation=True`` a lax coercion like ``{"n": "1"}`` for + an ``int`` parameter is rejected on the synchronous path. Task submission must + reject it identically rather than silently coercing and queueing it. """ + mcp = FastMCP("strict-task-server", strict_input_validation=True) + mcp.add_extension(TasksExtension()) - class _Recorder(MessageHandler): - def __init__(self): - super().__init__() - self.methods: list[str] = [] - - async def on_notification(self, message: mcp_types.ServerNotification) -> None: - self.methods.append(message.method) - - server = FastMCP("strict-task-server", strict_input_validation=True) - - @server.tool(task=True) + @mcp.tool(task=True) async def square(n: int) -> int: return n * n - recorder = _Recorder() - async with Client(server, mode="legacy", message_handler=recorder) as client: + async with running_task_server(mcp): # Sync path rejects the string-for-int coercion under strict validation. - with pytest.raises(ToolError): - await client.call_tool("square", {"n": "1"}) - - # Task path must reject it too, before any task state is created — so no - # status notification is emitted for the orphaned submission. - task = await client.call_tool("square", {"n": "1"}, task=True) - assert task.returned_immediately - with pytest.raises(ToolError): - await task.result() - - assert "notifications/tasks/status" not in recorder.methods + with pytest.raises(ValidationError): + await call_tool_without_optin(mcp, "square", {"n": "1"}) + # Task submission must reject it too, before any task state is created. + with pytest.raises(ValidationError): + await submit_task(mcp, "square", {"n": "1"}) -async def test_task_submission_valid_argument_under_strict_validation(): +async def test_valid_argument_submits_under_strict_validation(): """A well-typed argument still submits fine when strict validation is on.""" - server = FastMCP("strict-task-valid-server", strict_input_validation=True) + mcp = FastMCP("strict-task-valid-server", strict_input_validation=True) + mcp.add_extension(TasksExtension()) - @server.tool(task=True) + @mcp.tool(task=True) async def square(n: int) -> int: return n * n - async with Client(server, mode="legacy") as client: - task = await client.call_tool("square", {"n": 4}, task=True) - assert not task.returned_immediately - result = await task.result() - assert result.data == 16 + async with running_task_server(mcp): + final = await run_task(mcp, "square", {"n": 4}) + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": 16} def test_resolve_param_hints_handles_partials(): @@ -176,72 +156,59 @@ def test_resolve_param_hints_handles_partials(): assert hints["items"] == list[_Item] -async def test_synchronous_tool_call_unchanged(tool_server): - """Tools without task metadata execute synchronously as before.""" - async with Client(tool_server, mode="legacy") as client: - # Regular call without task metadata - result = await client.call_tool("simple_tool", {"message": "hello"}) - - # Should execute immediately and return result - assert "Processed: hello" in str(result) +# --------------------------------------------------------------------------- +# Decoration and execution +# --------------------------------------------------------------------------- -async def test_tool_with_task_metadata_returns_immediately(tool_server): - """Tools with task metadata return immediately with ToolTask object.""" - async with Client(tool_server, mode="legacy") as client: - # Call with task metadata - task = await client.call_tool("simple_tool", {"message": "test"}, task=True) - assert task - assert not task.returned_immediately - - assert isinstance(task, ToolTask) - assert isinstance(task.task_id, str) - assert len(task.task_id) > 0 +async def test_synchronous_tool_call_without_opt_in(): + """A tool called without a tasks opt-in executes synchronously as before.""" + mcp = _tool_server() + async with running_task_server(mcp): + result = await call_tool_without_optin(mcp, "simple_tool", {"message": "hello"}) + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "Processed: hello"} -async def test_tool_task_executes_in_background(tool_server): - """Tool task is submitted to Docket and executes in background.""" - execution_started = asyncio.Event() - execution_completed = asyncio.Event() +async def test_tool_task_returns_metadata_immediately(): + """Submitting a task returns task metadata with a server-generated id.""" + mcp = _tool_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "simple_tool", {"message": "test"}) + assert isinstance(created, CreateTaskResult) + assert isinstance(created.task_id, str) + assert created.task_id + assert created.status == "working" - @tool_server.tool(task=True) - async def coordinated_tool() -> str: - """Tool with coordination points.""" - execution_started.set() - await execution_completed.wait() + +async def test_tool_task_executes_in_background(): + """A submitted task runs in the background and can be polled to completion.""" + mcp = FastMCP("bg-server") + mcp.add_extension(TasksExtension()) + started = asyncio.Event() + finish = asyncio.Event() + + @mcp.tool(task=True) + async def coordinated() -> str: + started.set() + await finish.wait() return "completed" - async with Client(tool_server, mode="legacy") as client: - task = await client.call_tool("coordinated_tool", task=True) - assert task - assert not task.returned_immediately - - # Wait for execution to start - await asyncio.wait_for(execution_started.wait(), timeout=2.0) - - # Task should still be working - status = await task.status() - assert status.status in ["working"] - - # Signal completion - execution_completed.set() - await task.wait(timeout=2.0) - - result = await task.result() - assert result.data == "completed" + async with running_task_server(mcp): + created = await submit_task(mcp, "coordinated", {}) + await asyncio.wait_for(started.wait(), timeout=2.0) + working = await get_task(mcp, created.task_id) + assert working.status == "working" + finish.set() + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.result["structuredContent"] == {"result": "completed"} -async def test_forbidden_mode_tool_rejects_task_calls(tool_server): - """Tools with task=False (mode=forbidden) reject task-augmented calls.""" - async with Client(tool_server, mode="legacy") as client: - # Calling with task=True when task=False should return error - task = await client.call_tool( - "sync_only_tool", {"message": "test"}, task=True, raise_on_error=False - ) - assert task - assert task.returned_immediately - - result = await task.result() - # New behavior: mode="forbidden" returns an error - assert result.is_error - assert "does not support task-augmented execution" in str(result) +async def test_forbidden_tool_runs_sync_even_with_opt_in(): + """A task=False tool runs synchronously even when the caller opts into tasks.""" + mcp = _tool_server() + async with running_task_server(mcp): + result = await _opted_in_call(mcp, "sync_only_tool", {"message": "test"}) + assert not isinstance(result, CreateTaskResult) + assert result.structured_content == {"result": "Sync: test"} diff --git a/tests/tasks/server/test_task_ttl.py b/tests/tasks/server/test_task_ttl.py index 4464f8de3..a3bb74b2d 100644 --- a/tests/tasks/server/test_task_ttl.py +++ b/tests/tasks/server/test_task_ttl.py @@ -1,26 +1,30 @@ -""" -Tests for SEP-1686 ttl parameter handling. +"""TTL handling for SEP-2663 tasks. -Per the spec, servers MUST return ttl in all tasks/get responses, -and results should be retained for ttl milliseconds after completion. +Servers report `ttlMs` in the create result and in every `tasks/get` response — +while the task is working and after it completes — using Docket's default +execution TTL (900000 ms) when none is configured. """ +from __future__ import annotations + import asyncio -import pytest - from fastmcp import FastMCP -from fastmcp.client import Client - -pytestmark = pytest.mark.skip( - reason="Phase 3: requires TasksExtension (SEP-2663 adapter)" +from fastmcp_tasks import TasksExtension +from tests.tasks.task_helpers import ( + get_task, + running_task_server, + submit_task, + wait_for_task, ) +# Docket's default execution_ttl is 900 seconds. +DEFAULT_TTL_MS = 900000 -@pytest.fixture -async def keepalive_server(): - """Create a server for testing ttl behavior.""" + +def _ttl_server() -> FastMCP: mcp = FastMCP("keepalive-test") + mcp.add_extension(TasksExtension()) @mcp.tool(task=True) async def quick_task(value: int) -> int: @@ -28,67 +32,40 @@ async def keepalive_server(): @mcp.tool(task=True) async def slow_task() -> str: - # Never completes during the test - the only test that submits this - # task checks status immediately after submission and never awaits - # completion, so there's no need for a real-time sleep here. + # Never completes during the test; the test only checks status/TTL while + # the task is still working, so a suspended coroutine is enough. await asyncio.Event().wait() return "done" return mcp -async def test_keepalive_returned_in_submitted_state(keepalive_server: FastMCP): - """ttl is returned in tasks/get even when task is submitted/working.""" - async with Client(keepalive_server, mode="legacy") as client: - # Submit task with explicit ttl - task = await client.call_tool( - "slow_task", - {}, - task=True, - ttl=30000, # 30 seconds (client-requested) - ) - - # Check status immediately - should be submitted or working - status = await task.status() - assert status.status in ["working"] - - # ttl should be present per spec (MUST return in all responses) - # TODO: Docket uses a global execution_ttl for all tasks, not per-task TTLs. - # The spec allows servers to override client-requested TTL (line 431). - # FastMCP returns the server's actual global TTL (60000ms default from Docket). - # If Docket gains per-task TTL support, update this to verify client-requested TTL is respected. - assert status.ttl == 60000 # Server's global TTL, not client-requested 30000 +async def test_ttl_returned_while_working(): + """ttlMs is present in the create result and in tasks/get while working.""" + mcp = _ttl_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "slow_task", {}) + assert created.ttl_ms == DEFAULT_TTL_MS + got = await get_task(mcp, created.task_id) + assert got.status == "working" + assert got.ttl_ms == DEFAULT_TTL_MS -async def test_keepalive_returned_in_completed_state(keepalive_server: FastMCP): - """ttl is returned in tasks/get after task completes.""" - async with Client(keepalive_server, mode="legacy") as client: - # Submit and complete task - task = await client.call_tool( - "quick_task", - {"value": 5}, - task=True, - ttl=45000, # Client-requested TTL - ) - await task.wait(timeout=2.0) - - # Check status - should be completed - status = await task.status() - assert status.status == "completed" - - # TODO: Docket uses global execution_ttl, not per-task TTLs. - # Server returns its global TTL (60000ms), not the client-requested 45000ms. - # This is spec-compliant - servers MAY override requested TTL (spec line 431). - assert status.ttl == 60000 # Server's global TTL, not client-requested 45000 +async def test_ttl_returned_after_completion(): + """ttlMs is present in tasks/get after the task completes.""" + mcp = _ttl_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "quick_task", {"value": 5}) + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + assert final.ttl_ms == DEFAULT_TTL_MS -async def test_default_keepalive_when_not_specified(keepalive_server: FastMCP): - """Default ttl is used when client doesn't specify.""" - async with Client(keepalive_server, mode="legacy") as client: - # Submit without explicit ttl - task = await client.call_tool("quick_task", {"value": 3}, task=True) - await task.wait(timeout=2.0) - - status = await task.status() - # Should have default ttl (60000ms = 60 seconds) - assert status.ttl == 60000 +async def test_default_ttl_when_unspecified(): + """The server applies Docket's default TTL when none is configured.""" + mcp = _ttl_server() + async with running_task_server(mcp): + created = await submit_task(mcp, "quick_task", {"value": 3}) + assert created.ttl_ms == DEFAULT_TTL_MS + got = await get_task(mcp, created.task_id) + assert got.ttl_ms == DEFAULT_TTL_MS diff --git a/tests/tasks/server/test_wire_models.py b/tests/tasks/server/test_wire_models.py new file mode 100644 index 000000000..b36486e97 --- /dev/null +++ b/tests/tasks/server/test_wire_models.py @@ -0,0 +1,123 @@ +"""Validate the SEP-2663 wire models against the vendored draft JSON schema. + +The models in `fastmcp_tasks.models` serialize to the `io.modelcontextprotocol/tasks` +extension shapes. This suite validates a serialized instance of each result shape +against the corresponding `$defs` entry in the vendored draft schema +(`tests/fixtures/ext-tasks-schema-draft.json`), so wire drift is caught here. + +The vendored schema composes results as `allOf[Result, Task]` where the Task arm +carries `additionalProperties: false`; a stray `_meta` therefore fails +validation. The models omit `_meta` and the runner's `exclude_none` dump keeps it +out, which is exactly what these assertions check. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from fastmcp_tasks.models import ( + CancelTaskResult, + CreateTaskResult, + GetTaskResult, + UpdateTaskResult, +) +from jsonschema import Draft202012Validator + +_SCHEMA = json.loads( + (Path(__file__).parents[2] / "fixtures" / "ext-tasks-schema-draft.json").read_text() +) +_DEFS = _SCHEMA["$defs"] + +_ISO = "2026-07-21T12:00:00+00:00" + + +def _validate(def_name: str, instance: dict[str, Any]) -> None: + schema = {"$defs": _DEFS, **_DEFS[def_name]} + Draft202012Validator(schema).validate(instance) + + +def _dump(model: Any) -> dict[str, Any]: + return model.model_dump(by_alias=True, mode="json", exclude_none=True) + + +def test_create_task_result_matches_schema(): + result = CreateTaskResult( + task_id="t1", + status="working", + created_at=_ISO, + last_updated_at=_ISO, + ttl_ms=900000, + poll_interval_ms=5000, + ) + _validate("CreateTaskResult", _dump(result)) + + +@pytest.mark.parametrize( + ("status", "payload"), + [ + ("working", {}), + ("completed", {"result": {"content": [], "isError": False}}), + ("failed", {"error": {"code": -32603, "message": "boom"}}), + ( + "input_required", + { + "input_requests": { + "k1": {"method": "elicitation/create", "params": {"message": "?"}} + } + }, + ), + ("cancelled", {}), + ], +) +def test_get_task_result_matches_schema(status: str, payload: dict[str, Any]): + result = GetTaskResult( + task_id="t1", + status=status, # type: ignore[arg-type] + created_at=_ISO, + last_updated_at=_ISO, + ttl_ms=900000, + poll_interval_ms=5000, + **payload, + ) + _validate("GetTaskResult", _dump(result)) + + +def test_get_task_result_completed_omits_error_and_inputs(): + """A completed result carries only `result` (the union arm forbids the rest).""" + result = GetTaskResult( + task_id="t1", + status="completed", + created_at=_ISO, + last_updated_at=_ISO, + ttl_ms=900000, + result={"content": [], "isError": False}, + ) + dumped = _dump(result) + assert "error" not in dumped + assert "inputRequests" not in dumped + + +def test_null_ttl_is_permitted_by_schema(): + """`ttlMs` is required-but-nullable; a null TTL still validates.""" + result = CreateTaskResult( + task_id="t1", + status="working", + created_at=_ISO, + last_updated_at=_ISO, + ttl_ms=None, + ) + dumped = result.model_dump(by_alias=True, mode="json", exclude_none=False) + # Drop the other None optionals the runner would also drop, keeping ttlMs=null. + dumped = { + k: v for k, v in dumped.items() if v is not None or k == "ttlMs" + } + _validate("CreateTaskResult", dumped) + + +@pytest.mark.parametrize("model", [UpdateTaskResult(), CancelTaskResult()]) +def test_ack_results_match_schema(model: Any): + def_name = type(model).__name__ + _validate(def_name, _dump(model)) diff --git a/tests/tasks/task_helpers.py b/tests/tasks/task_helpers.py new file mode 100644 index 000000000..3f8563b12 --- /dev/null +++ b/tests/tasks/task_helpers.py @@ -0,0 +1,210 @@ +"""Shared helpers for driving SEP-2663 tasks in server-side tests. + +There is no client task-submission API until Phase 4, so server-side tests drive +the task lifecycle in-process: the create decision runs through the real +`tools/call` interceptor (with a per-request tasks opt-in bound into the request +context), and `tasks/get` / `tasks/update` / `tasks/cancel` call the extension's +handler functions directly. Optional auth binding exercises the auth-scoped task +isolation. + +Typical use:: + + async with running_task_server(mcp): + created = await submit_task(mcp, "square", {"n": 6}) + final = await wait_for_task(mcp, created.task_id) + assert final.status == "completed" + +or the one-shot:: + + async with running_task_server(mcp): + final = await run_task(mcp, "square", {"n": 6}) +""" + +from __future__ import annotations + +import asyncio +import contextlib +from types import SimpleNamespace +from typing import Any + +from fastmcp_tasks.handlers import tasks_cancel, tasks_get, tasks_update +from fastmcp_tasks.models import ( + CancelTaskResult, + CreateTaskResult, + GetTaskResult, + UpdateTaskResult, +) +from mcp.server.auth.middleware.auth_context import auth_context_var +from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser +from mcp.server.context import ServerRequestContext +from mcp_types import CLIENT_CAPABILITIES_META_KEY + +from fastmcp.server.auth import AccessToken +from fastmcp.server.dependencies import bind_request_context +from fastmcp.server.server import FastMCP +from fastmcp.utilities.tasks import TASKS_EXTENSION_ID + +TERMINAL_STATES = frozenset({"completed", "failed", "cancelled"}) + + +def opt_in_meta(settings: dict[str, Any] | None = None) -> dict[str, Any]: + """The per-request `_meta` block that opts the tasks extension in.""" + return { + CLIENT_CAPABILITIES_META_KEY: { + "extensions": {TASKS_EXTENSION_ID: settings or {}} + } + } + + +def make_access_token(client_id: str, sub: str | None = None) -> AccessToken: + """A minimal FastMCP access token for auth-scoped task tests.""" + claims: dict[str, Any] = {"sub": sub} if sub is not None else {} + return AccessToken( + token=f"token-{client_id}-{sub}", + client_id=client_id, + scopes=[], + claims=claims, + ) + + +@contextlib.contextmanager +def auth_scope(access_token: AccessToken | None): + """Bind (or clear) the auth context so `get_task_scope` sees a caller.""" + if access_token is None: + yield + return + token = auth_context_var.set(AuthenticatedUser(access_token)) + try: + yield + finally: + auth_context_var.reset(token) + + +@contextlib.contextmanager +def _opted_in_request( + name: str, arguments: dict[str, Any] | None, settings: dict[str, Any] | None +): + """Bind a request context carrying the tasks opt-in for a `tools/call`.""" + params: dict[str, Any] = { + "name": name, + "arguments": arguments or {}, + "_meta": opt_in_meta(settings), + } + srctx = ServerRequestContext( + session=SimpleNamespace(), + lifespan_context={}, + protocol_version="2026-07-28", + method="tools/call", + params=params, + ) + with bind_request_context(srctx): + yield + + +def running_task_server(server: FastMCP): + """Enter the server lifespan (Docket backend + worker) for the block.""" + return server._lifespan_manager() + + +async def submit_task( + server: FastMCP, + name: str, + arguments: dict[str, Any] | None = None, + *, + access_token: AccessToken | None = None, + settings: dict[str, Any] | None = None, +) -> CreateTaskResult: + """Run an opted-in `tools/call` through the interceptor and return its task.""" + with auth_scope(access_token), _opted_in_request(name, arguments, settings): + result = await server.call_tool(name, arguments or {}) + if not isinstance(result, CreateTaskResult): + raise AssertionError( + f"Expected the call to be tasked, got {type(result).__name__}: {result!r}" + ) + return result + + +async def call_tool_without_optin( + server: FastMCP, + name: str, + arguments: dict[str, Any] | None = None, + *, + access_token: AccessToken | None = None, +): + """Run a `tools/call` with no tasks opt-in (synchronous unless mode=required).""" + with auth_scope(access_token): + return await server.call_tool(name, arguments or {}) + + +async def get_task( + server: FastMCP, + task_id: str, + *, + access_token: AccessToken | None = None, +) -> GetTaskResult: + """Call the `tasks/get` handler within the given auth scope.""" + with auth_scope(access_token): + return await tasks_get(server, task_id) + + +async def update_task( + server: FastMCP, + task_id: str, + input_responses: dict[str, Any], + *, + access_token: AccessToken | None = None, +) -> UpdateTaskResult: + """Call the `tasks/update` handler within the given auth scope.""" + with auth_scope(access_token): + return await tasks_update(server, task_id, input_responses) + + +async def cancel_task( + server: FastMCP, + task_id: str, + *, + access_token: AccessToken | None = None, +) -> CancelTaskResult: + """Call the `tasks/cancel` handler within the given auth scope.""" + with auth_scope(access_token): + return await tasks_cancel(server, task_id) + + +async def wait_for_task( + server: FastMCP, + task_id: str, + *, + access_token: AccessToken | None = None, + target_states: frozenset[str] = TERMINAL_STATES, + timeout: float = 5.0, + poll: float = 0.02, +) -> GetTaskResult: + """Poll `tasks/get` until the task reaches one of `target_states`.""" + deadline = asyncio.get_event_loop().time() + timeout + result = await get_task(server, task_id, access_token=access_token) + while result.status not in target_states: + if asyncio.get_event_loop().time() >= deadline: + raise TimeoutError( + f"Task {task_id} still {result.status!r} after {timeout}s " + f"(waiting for {sorted(target_states)})" + ) + await asyncio.sleep(poll) + result = await get_task(server, task_id, access_token=access_token) + return result + + +async def run_task( + server: FastMCP, + name: str, + arguments: dict[str, Any] | None = None, + *, + access_token: AccessToken | None = None, + timeout: float = 5.0, +) -> GetTaskResult: + """Submit a task and wait for it to reach a terminal state.""" + created = await submit_task( + server, name, arguments, access_token=access_token + ) + return await wait_for_task( + server, created.task_id, access_token=access_token, timeout=timeout + )