diff --git a/examples/providers/sqlite/server.py b/examples/providers/sqlite/server.py index 3caa8bde0..2d4c39893 100644 --- a/examples/providers/sqlite/server.py +++ b/examples/providers/sqlite/server.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio import json +from collections.abc import Sequence from pathlib import Path from typing import Any @@ -21,7 +22,6 @@ import aiosqlite from rich import print from fastmcp import Client, FastMCP -from fastmcp.server.context import Context from fastmcp.server.providers import Provider from fastmcp.tools.tool import Tool, ToolResult @@ -72,16 +72,17 @@ class SQLiteToolProvider(Provider): """ def __init__(self, db_path: str): + super().__init__() self.db_path = db_path - async def list_tools(self, context: Context) -> list[Tool]: + async def list_tools(self) -> Sequence[Tool]: async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute("SELECT * FROM tools WHERE enabled = 1") as cursor: rows = await cursor.fetchall() return [self._make_tool(row) for row in rows] - async def get_tool(self, context: Context, name: str) -> Tool | None: + async def get_tool(self, name: str) -> Tool | None: async with aiosqlite.connect(self.db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( diff --git a/src/fastmcp/prompts/prompt.py b/src/fastmcp/prompts/prompt.py index 394a5f714..02f88df0a 100644 --- a/src/fastmcp/prompts/prompt.py +++ b/src/fastmcp/prompts/prompt.py @@ -13,6 +13,7 @@ import pydantic_core if TYPE_CHECKING: from docket import Docket from docket.execution import Execution +import mcp.types from mcp import GetPromptResult from mcp.types import ContentBlock, Icon, PromptMessage, Role, TextContent from mcp.types import Prompt as SDKPrompt @@ -211,19 +212,70 @@ class Prompt(FastMCPComponent): def convert_result(self, raw_value: Any) -> PromptResult: """Convert a raw return value to PromptResult. - Subclasses should override this to handle their specific conversion logic. + Handles PromptResult passthrough and converts raw values to messages. """ - raise NotImplementedError("Subclasses must implement convert_result()") + if isinstance(raw_value, PromptResult): + return raw_value + + # Normalize to list + if not isinstance(raw_value, list | tuple): + raw_value = [raw_value] + + # Convert result to messages + messages: list[PromptMessage] = [] + for msg in raw_value: + try: + if isinstance(msg, PromptMessage): + messages.append(msg) + elif isinstance(msg, str): + messages.append( + PromptMessage( + role="user", + content=TextContent(type="text", text=msg), + ) + ) + else: + content = pydantic_core.to_json(msg, fallback=str).decode() + messages.append( + PromptMessage( + role="user", + content=TextContent(type="text", text=content), + ) + ) + except Exception as e: + raise PromptError("Could not convert prompt result to message.") from e + + return PromptResult( + messages=messages, + description=self.description, + meta=self.meta, + ) async def _render( self, arguments: dict[str, Any] | None = None, - ) -> PromptResult: - """Internal API that always returns PromptResult. + ) -> PromptResult | mcp.types.CreateTaskResult: + """Server entry point that handles task routing. - Calls render() and wraps list[PromptMessage] in PromptResult. - This is what PromptManager calls internally. + This allows ANY Prompt subclass to support background execution by setting + task_config.mode to "supported" or "required". The server calls this + method instead of render() directly. + + Subclasses can override this to customize task routing behavior. + For example, FastMCPProviderPrompt overrides to delegate to child + middleware without submitting to Docket. """ + from fastmcp.server.dependencies import _docket_fn_key + from fastmcp.server.tasks.routing import check_background_task + + key = _docket_fn_key.get() or self.key + task_result = await check_background_task( + component=self, task_type="prompt", key=key, arguments=arguments + ) + if task_result: + return task_result + + # Synchronous execution result = await self.render(arguments) if isinstance(result, PromptResult): return result @@ -247,10 +299,27 @@ class Prompt(FastMCPComponent): docket.register(self.render, names=[self.key]) async def add_to_docket( # type: ignore[override] - self, docket: Docket, arguments: dict[str, Any] | None, **kwargs: Any + self, + docket: Docket, + arguments: dict[str, Any] | None, + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: - """Schedule this prompt for background execution via docket.""" - return await docket.add(self.key, **kwargs)(arguments) + """Schedule this prompt for background execution via docket. + + Args: + docket: The Docket instance + arguments: Prompt arguments + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() + """ + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(arguments) class FunctionPrompt(Prompt): @@ -444,46 +513,6 @@ class FunctionPrompt(Prompt): logger.exception(f"Error rendering prompt {self.name}") raise PromptError(f"Error rendering prompt {self.name}.") from e - def convert_result(self, raw_value: Any) -> PromptResult: - """Convert a raw return value to PromptResult. - - This handles the same conversion logic as render(), but works on - already-executed raw values (e.g., from Docket background execution). - """ - # Normalize to list - if not isinstance(raw_value, list | tuple): - raw_value = [raw_value] - - # Convert result to messages - messages: list[PromptMessage] = [] - for msg in raw_value: - try: - if isinstance(msg, PromptMessage): - messages.append(msg) - elif isinstance(msg, str): - messages.append( - PromptMessage( - role="user", - content=TextContent(type="text", text=msg), - ) - ) - else: - content = pydantic_core.to_json(msg, fallback=str).decode() - messages.append( - PromptMessage( - role="user", - content=TextContent(type="text", text=content), - ) - ) - except Exception as e: - raise PromptError("Could not convert prompt result to message.") from e - - return PromptResult( - messages=messages, - description=self.description, - meta=self.meta, - ) - def register_with_docket(self, docket: Docket) -> None: """Register this prompt with docket for background execution. @@ -495,10 +524,26 @@ class FunctionPrompt(Prompt): docket.register(self.fn, names=[self.key]) # type: ignore[arg-type] async def add_to_docket( # type: ignore[override] - self, docket: Docket, arguments: dict[str, Any] | None, **kwargs: Any + self, + docket: Docket, + arguments: dict[str, Any] | None, + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: """Schedule this prompt for background execution via docket. FunctionPrompt splats the arguments dict since .fn expects **kwargs. + + Args: + docket: The Docket instance + arguments: Prompt arguments + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() """ - return await docket.add(self.key, **kwargs)(**(arguments or {})) + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(**(arguments or {})) diff --git a/src/fastmcp/prompts/prompt_manager.py b/src/fastmcp/prompts/prompt_manager.py index 6dd57a8a9..8b19ded92 100644 --- a/src/fastmcp/prompts/prompt_manager.py +++ b/src/fastmcp/prompts/prompt_manager.py @@ -4,6 +4,8 @@ import warnings from collections.abc import Awaitable, Callable from typing import Any +import mcp.types + from fastmcp import settings from fastmcp.exceptions import NotFoundError, PromptError from fastmcp.prompts.prompt import ( @@ -105,12 +107,16 @@ class PromptManager: self, name: str, arguments: dict[str, Any] | None = None, - ) -> PromptResult: + ) -> PromptResult | mcp.types.CreateTaskResult: """ Internal API for servers: Finds and renders a prompt. Note: Full error handling (logging, masking) is done at the FastMCP server level. This method provides basic error wrapping for direct usage. + + Returns: + PromptResult for synchronous execution, or CreateTaskResult if + the prompt was submitted to Docket for background execution. """ prompt = await self.get_prompt(name) try: diff --git a/src/fastmcp/resources/resource.py b/src/fastmcp/resources/resource.py index 60e6cd3ae..d6a98b1b6 100644 --- a/src/fastmcp/resources/resource.py +++ b/src/fastmcp/resources/resource.py @@ -231,13 +231,35 @@ class Resource(FastMCPComponent): """ raise NotImplementedError("Subclasses must implement read()") - async def _read(self) -> ResourceContent: - """Internal API that always returns ResourceContent. + def convert_result(self, raw_value: Any) -> ResourceContent: + """Convert a raw return value to ResourceContent. - This method calls read() and wraps str/bytes results in ResourceContent. - ResourceManager and other internal code should call this method instead - of read() directly. + Handles ResourceContent passthrough and converts raw values using mime_type. """ + return ResourceContent.from_value(raw_value, mime_type=self.mime_type) + + async def _read(self) -> ResourceContent | mcp.types.CreateTaskResult: + """Server entry point that handles task routing. + + This allows ANY Resource subclass to support background execution by setting + task_config.mode to "supported" or "required". The server calls this + method instead of read() directly. + + Subclasses can override this to customize task routing behavior. + For example, FastMCPProviderResource overrides to delegate to child + middleware without submitting to Docket. + """ + from fastmcp.server.dependencies import _docket_fn_key + from fastmcp.server.tasks.routing import check_background_task + + key = _docket_fn_key.get() or self.key + task_result = await check_background_task( + component=self, task_type="resource", key=key + ) + if task_result: + return task_result + + # Synchronous execution result = await self.read() if isinstance(result, ResourceContent): return result @@ -250,7 +272,7 @@ class Resource(FastMCPComponent): DeprecationWarning, stacklevel=2, ) - return ResourceContent.from_value(result, mime_type=self.mime_type) + return self.convert_result(result) def to_mcp_resource( self, @@ -288,10 +310,25 @@ class Resource(FastMCPComponent): docket.register(self.read, names=[self.key]) async def add_to_docket( # type: ignore[override] - self, docket: Docket, **kwargs: Any + self, + docket: Docket, + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: - """Schedule this resource for background execution via docket.""" - return await docket.add(self.key, **kwargs)() + """Schedule this resource for background execution via docket. + + Args: + docket: The Docket instance + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() + """ + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)() class FunctionResource(Resource): @@ -378,14 +415,6 @@ class FunctionResource(Resource): return self.convert_result(result) - def convert_result(self, raw_value: Any) -> ResourceContent: - """Convert a raw return value to ResourceContent. - - This handles the same conversion logic as read(), but works on - already-executed raw values (e.g., from Docket background execution). - """ - return ResourceContent.from_value(raw_value, mime_type=self.mime_type) - def register_with_docket(self, docket: Docket) -> None: """Register this resource with docket for background execution. diff --git a/src/fastmcp/resources/resource_manager.py b/src/fastmcp/resources/resource_manager.py index aea3f1d09..ecd7b06a3 100644 --- a/src/fastmcp/resources/resource_manager.py +++ b/src/fastmcp/resources/resource_manager.py @@ -7,6 +7,7 @@ import warnings from collections.abc import Callable from typing import Any +import mcp.types from pydantic import AnyUrl from fastmcp import settings @@ -286,7 +287,9 @@ class ResourceManager: raise NotFoundError(f"Unknown resource: {uri_str}") - async def read_resource(self, uri: AnyUrl | str) -> ResourceContent: + async def read_resource( + self, uri: AnyUrl | str + ) -> ResourceContent | mcp.types.CreateTaskResult: """ Internal API for servers: Finds and reads a resource. @@ -294,7 +297,8 @@ class ResourceManager: server level. This method provides basic error wrapping for direct usage. Returns: - ResourceContent: The canonical content wrapper. + ResourceContent for synchronous execution, or CreateTaskResult if + the resource was submitted to Docket for background execution. """ uri_str = str(uri) diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index 19a8722d9..c03258e7b 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -8,6 +8,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any from urllib.parse import parse_qs, unquote +import mcp.types from mcp.types import Annotations, Icon if TYPE_CHECKING: @@ -20,7 +21,7 @@ from pydantic import ( validate_call, ) -from fastmcp.resources.resource import Resource +from fastmcp.resources.resource import Resource, ResourceContent from fastmcp.server.dependencies import get_context, without_injected_parameters from fastmcp.server.tasks.config import TaskConfig from fastmcp.utilities.components import FastMCPComponent @@ -176,6 +177,47 @@ class ResourceTemplate(FastMCPComponent): "Subclasses must implement read() or override create_resource()" ) + def convert_result(self, raw_value: Any) -> ResourceContent: + """Convert a raw return value to ResourceContent. + + Handles ResourceContent passthrough and converts raw values using mime_type. + """ + return ResourceContent.from_value(raw_value, mime_type=self.mime_type) + + async def _read( + self, uri: str, params: dict[str, Any] + ) -> ResourceContent | mcp.types.CreateTaskResult: + """Server entry point that handles task routing. + + This allows ANY ResourceTemplate subclass to support background execution + by setting task_config.mode to "supported" or "required". The server calls + this method instead of create_resource()/read() directly. + + Subclasses can override this to customize task routing behavior. + For example, FastMCPProviderResourceTemplate overrides to delegate to child + middleware without submitting to Docket. + """ + from fastmcp.server.dependencies import _docket_fn_key + from fastmcp.server.tasks.routing import check_background_task + + # Templates need pattern check: only use contextvar if it contains '{' + key = _docket_fn_key.get() + if not key or "{" not in key: + key = self.key + task_result = await check_background_task( + component=self, task_type="template", key=key, arguments=params + ) + if task_result: + return task_result + + # Synchronous execution - create resource and read directly + # Call resource.read() not resource._read() to avoid task routing on ephemeral resource + resource = await self.create_resource(uri, params) + result = await resource.read() + if isinstance(result, ResourceContent): + return result + return resource.convert_result(result) + async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: """Create a resource from the template with the given parameters. @@ -233,10 +275,27 @@ class ResourceTemplate(FastMCPComponent): docket.register(self.read, names=[self.key]) async def add_to_docket( # type: ignore[override] - self, docket: Docket, params: dict[str, Any], **kwargs: Any + self, + docket: Docket, + params: dict[str, Any], + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: - """Schedule this template for background execution via docket.""" - return await docket.add(self.key, **kwargs)(params) + """Schedule this template for background execution via docket. + + Args: + docket: The Docket instance + params: Template parameters + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() + """ + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(params) class FunctionResourceTemplate(ResourceTemplate): @@ -244,6 +303,31 @@ class FunctionResourceTemplate(ResourceTemplate): fn: Callable[..., Any] + async def _read( + self, uri: str, params: dict[str, Any] + ) -> ResourceContent | mcp.types.CreateTaskResult: + """Optimized server entry point that skips ephemeral resource creation. + + For FunctionResourceTemplate, we can call read() directly instead of + creating a temporary resource, which is more efficient. + """ + from fastmcp.server.dependencies import _docket_fn_key + from fastmcp.server.tasks.routing import check_background_task + + # Templates need pattern check: only use contextvar if it contains '{' + key = _docket_fn_key.get() + if not key or "{" not in key: + key = self.key + task_result = await check_background_task( + component=self, task_type="template", key=key, arguments=params + ) + if task_result: + return task_result + + # Synchronous execution - call read() directly, skip resource creation + result = await self.read(arguments=params) + return self.convert_result(result) + async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: """Create a resource from the template with the given parameters.""" @@ -305,13 +389,29 @@ class FunctionResourceTemplate(ResourceTemplate): docket.register(self.fn, names=[self.key]) async def add_to_docket( # type: ignore[override] - self, docket: Docket, params: dict[str, Any], **kwargs: Any + self, + docket: Docket, + params: dict[str, Any], + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: """Schedule this template for background execution via docket. FunctionResourceTemplate splats the params dict since .fn expects **kwargs. + + Args: + docket: The Docket instance + params: Template parameters + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() """ - return await docket.add(self.key, **kwargs)(**params) + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(**params) @classmethod def from_function( diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index 3671ba42b..545c1f019 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -40,6 +40,16 @@ _current_server: ContextVar[weakref.ref[FastMCP] | None] = ContextVar( # type: "server", default=None ) +# ContextVar for propagating task metadata through the async call stack +# When set, indicates this call should be executed as a background task +_task_metadata: ContextVar[dict[str, Any] | None] = ContextVar( + "task_metadata", default=None +) + +# ContextVar for the component's Docket function lookup key (with namespace prefix) +# Used by Tool._run(), Resource._read(), Prompt._render() to find the registered function +_docket_fn_key: ContextVar[str | None] = ContextVar("docket_fn_key", default=None) + __all__ = [ "AccessToken", "CurrentContext", @@ -52,6 +62,7 @@ __all__ = [ "get_http_headers", "get_http_request", "get_server", + "get_task_metadata", "resolve_dependencies", "without_injected_parameters", ] @@ -266,6 +277,16 @@ def get_context() -> Context: return context +def get_task_metadata() -> dict[str, Any] | None: + """Get the current task metadata from the context. + + Returns: + The task metadata dict if this is a background task request, + or None if this is a normal execution. + """ + return _task_metadata.get() + + class _CurrentContext(Dependency): """Internal dependency class for CurrentContext.""" diff --git a/src/fastmcp/server/providers/base.py b/src/fastmcp/server/providers/base.py index c15bc24ee..c1d1d3dba 100644 --- a/src/fastmcp/server/providers/base.py +++ b/src/fastmcp/server/providers/base.py @@ -31,18 +31,11 @@ from __future__ import annotations from collections.abc import AsyncIterator, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import TYPE_CHECKING, Any -from fastmcp.prompts.prompt import Prompt, PromptResult -from fastmcp.resources.resource import Resource, ResourceContent +from fastmcp.prompts.prompt import Prompt +from fastmcp.resources.resource import Resource from fastmcp.resources.template import ResourceTemplate -from fastmcp.tools.tool import Tool, ToolResult - -if TYPE_CHECKING: - from fastmcp.prompts.prompt import FunctionPrompt - from fastmcp.resources.resource import FunctionResource - from fastmcp.resources.template import FunctionResourceTemplate - from fastmcp.tools.tool import FunctionTool +from fastmcp.tools.tool import Tool @dataclass @@ -57,16 +50,16 @@ class Components: @dataclass class TaskComponents: - """Collection of function-based components eligible for background task execution. + """Collection of components eligible for background task execution. Used by get_tasks() to return components for Docket registration. - All components have a `.fn` attribute pointing to the underlying callable. + Components must implement register_with_docket() and add_to_docket(). """ - tools: Sequence[FunctionTool] = () - resources: Sequence[FunctionResource] = () - templates: Sequence[FunctionResourceTemplate] = () - prompts: Sequence[FunctionPrompt] = () + tools: Sequence[Tool] = () + resources: Sequence[Resource] = () + templates: Sequence[ResourceTemplate] = () + prompts: Sequence[Prompt] = () class Provider: @@ -80,13 +73,11 @@ class Provider: - Return `None` from `get_*` methods to indicate "I don't have it" (search continues) - Static components (registered via decorators) always take precedence over providers - Providers are queried in registration order; first non-None wins + - Components execute themselves via run()/read()/render() - providers just source them Error handling: - `list_*` methods: Errors are logged and the provider returns empty (graceful degradation). This allows other providers to still contribute their components. - - Execution methods (`call_tool`, `read_resource`, `render_prompt`): Errors propagate - with unified handling. ToolError/ResourceError/PromptError pass through; other - exceptions are wrapped with optional detail masking. """ def with_transforms( @@ -176,22 +167,6 @@ class Provider: tools = await self.list_tools() return next((t for t in tools if t.name == name), None) - async def call_tool( - self, name: str, arguments: dict[str, Any] - ) -> ToolResult | None: - """Execute a tool by name. - - Default implementation gets the tool and runs it. - Override for custom execution logic (e.g., middleware, error handling). - - Returns: - The ToolResult if found and executed, or None if tool not found. - """ - tool = await self.get_tool(name) - if tool is None: - return None - return await tool.run(arguments) - async def list_resources(self) -> Sequence[Resource]: """Return all available resources. @@ -234,43 +209,6 @@ class Provider: None, ) - async def read_resource(self, uri: str) -> ResourceContent | None: - """Read a concrete resource by URI. - - Default implementation gets the resource and reads it. - Override for custom read logic (e.g., middleware, caching). - - Note: This only handles concrete resources. For template-based resources, - use read_resource_template(). - - Returns: - The ResourceContent if found and read, or None if not found. - """ - resource = await self.get_resource(uri) - if resource is None: - return None - return await resource._read() - - async def read_resource_template(self, uri: str) -> ResourceContent | None: - """Read a resource via a matching template. - - Default implementation finds a matching template, creates a resource - from it, and reads the content. - Override for custom read logic (e.g., middleware, caching). - - Returns: - The ResourceContent if a matching template is found and read, - or None if no template matches. - """ - template = await self.get_resource_template(uri) - if template is None: - return None - params = template.matches(uri) - if params is None: - return None - resource = await template.create_resource(uri, params) - return await resource._read() - async def list_prompts(self) -> Sequence[Prompt]: """Return all available prompts. @@ -290,22 +228,6 @@ class Provider: prompts = await self.list_prompts() return next((p for p in prompts if p.name == name), None) - async def render_prompt( - self, name: str, arguments: dict[str, Any] | None - ) -> PromptResult | None: - """Render a prompt by name. - - Default implementation gets the prompt and renders it. - Override for custom render logic (e.g., middleware, templating). - - Returns: - The PromptResult if found and rendered, or None if not found. - """ - prompt = await self.get_prompt(name) - if prompt is None: - return None - return await prompt._render(arguments) - # ------------------------------------------------------------------------- # Task registration # ------------------------------------------------------------------------- diff --git a/src/fastmcp/server/providers/fastmcp_provider.py b/src/fastmcp/server/providers/fastmcp_provider.py index e1ca61855..06f3a415b 100644 --- a/src/fastmcp/server/providers/fastmcp_provider.py +++ b/src/fastmcp/server/providers/fastmcp_provider.py @@ -2,15 +2,22 @@ This module provides the `FastMCPProvider` class that wraps a FastMCP server and exposes its components through the Provider interface. + +It also provides FastMCPProvider* component classes that delegate execution to +the wrapped server's middleware, ensuring middleware runs when components are +executed. """ from __future__ import annotations +import re from collections.abc import AsyncIterator, Sequence from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any -from fastmcp.exceptions import NotFoundError +import mcp.types +from mcp.types import AnyUrl + from fastmcp.prompts.prompt import Prompt, PromptResult from fastmcp.resources.resource import Resource, ResourceContent from fastmcp.resources.template import ResourceTemplate @@ -18,11 +25,408 @@ from fastmcp.server.providers.base import Provider, TaskComponents from fastmcp.tools.tool import Tool, ToolResult if TYPE_CHECKING: - from fastmcp.prompts.prompt import FunctionPrompt - from fastmcp.resources.resource import FunctionResource - from fastmcp.resources.template import FunctionResourceTemplate + from docket import Docket + from docket.execution import Execution + from fastmcp.server.server import FastMCP - from fastmcp.tools.tool import FunctionTool + + +def _expand_uri_template(template: str, params: dict[str, Any]) -> str: + """Expand a URI template with parameters. + + Simple implementation that handles {name} style placeholders. + """ + result = template + for key, value in params.items(): + result = re.sub(rf"\{{{key}\}}", str(value), result) + return result + + +# ----------------------------------------------------------------------------- +# FastMCPProvider component classes +# ----------------------------------------------------------------------------- + + +class FastMCPProviderTool(Tool): + """Tool that delegates execution to a wrapped server's middleware. + + When `run()` is called, this tool invokes the wrapped server's + `_call_tool_middleware()` method, ensuring the server's middleware + chain is executed. + """ + + _server: Any = None # FastMCP, but Any to avoid circular import + _original_name: str | None = None + + def __init__( + self, + server: Any, + original_name: str, + **kwargs: Any, + ): + super().__init__(**kwargs) + self._server = server + self._original_name = original_name + + @classmethod + def wrap(cls, server: Any, tool: Tool) -> FastMCPProviderTool: + """Wrap a Tool to delegate execution to the server's middleware.""" + return cls( + server=server, + original_name=tool.name, + name=tool.name, + description=tool.description, + parameters=tool.parameters, + output_schema=tool.output_schema, + tags=tool.tags, + annotations=tool.annotations, + enabled=tool.enabled, + task_config=tool.task_config, + ) + + async def _run( + self, arguments: dict[str, Any] + ) -> ToolResult | mcp.types.CreateTaskResult: + """Skip task handling - delegate to run() which calls child middleware. + + The actual underlying tool will check _task_metadata contextvar and + submit to Docket if appropriate. This wrapper just passes through. + """ + return await self.run(arguments) + + async def run( + self, arguments: dict[str, Any] + ) -> ToolResult | mcp.types.CreateTaskResult: # type: ignore[override] + """Delegate to child server's middleware chain. + + This runs BEFORE any backgrounding decision - the actual underlying + tool will check contextvars and submit to Docket if appropriate. + """ + return await self._server._call_tool_middleware(self._original_name, arguments) + + +class FastMCPProviderResource(Resource): + """Resource that delegates reading to a wrapped server's middleware. + + When `read()` is called, this resource invokes the wrapped server's + `_read_resource_middleware()` method, ensuring the server's middleware + chain is executed. + """ + + _server: Any = None # FastMCP, but Any to avoid circular import + _original_uri: str | None = None + + def __init__( + self, + server: Any, + original_uri: str, + **kwargs: Any, + ): + super().__init__(**kwargs) + self._server = server + self._original_uri = original_uri + + @classmethod + def wrap(cls, server: Any, resource: Resource) -> FastMCPProviderResource: + """Wrap a Resource to delegate reading to the server's middleware.""" + return cls( + server=server, + original_uri=str(resource.uri), + uri=resource.uri, + name=resource.name, + description=resource.description, + mime_type=resource.mime_type, + tags=resource.tags, + annotations=resource.annotations, + enabled=resource.enabled, + task_config=resource.task_config, + ) + + async def _read(self) -> ResourceContent | mcp.types.CreateTaskResult: + """Skip task routing - delegate to read() which calls child middleware. + + The actual underlying resource will check _task_metadata contextvar and + submit to Docket if appropriate. This wrapper just passes through. + """ + return await self.read() + + async def read(self) -> ResourceContent | mcp.types.CreateTaskResult: # type: ignore[override] + """Delegate to child server's middleware. + + When called from a Docket worker (background task), there's no FastMCP + context set up, so we create one for the child server. + + Note: The _docket_fn_key contextvar is intentionally NOT updated here. + The parent set it to the full namespaced key (e.g., data://c/gc/value) + which is what the function is registered under in Docket. All provider + layers pass this through unchanged so the eventual resource._read() + uses the correct Docket lookup key. + """ + import fastmcp.server.context + + try: + from fastmcp.server.dependencies import get_context + + get_context() # Will raise if no context + result = await self._server._read_resource_middleware(self._original_uri) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result[0] + except RuntimeError: + # No context (e.g., Docket worker) - create one for the child server + async with fastmcp.server.context.Context(fastmcp=self._server): + result = await self._server._read_resource_middleware( + self._original_uri + ) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result[0] + + +class FastMCPProviderPrompt(Prompt): + """Prompt that delegates rendering to a wrapped server's middleware. + + When `render()` is called, this prompt invokes the wrapped server's + `_get_prompt_content_middleware()` method, ensuring the server's middleware + chain is executed. + """ + + _server: Any = None # FastMCP, but Any to avoid circular import + _original_name: str | None = None + + def __init__( + self, + server: Any, + original_name: str, + **kwargs: Any, + ): + super().__init__(**kwargs) + self._server = server + self._original_name = original_name + + @classmethod + def wrap(cls, server: Any, prompt: Prompt) -> FastMCPProviderPrompt: + """Wrap a Prompt to delegate rendering to the server's middleware.""" + return cls( + server=server, + original_name=prompt.name, + name=prompt.name, + description=prompt.description, + arguments=prompt.arguments, + tags=prompt.tags, + enabled=prompt.enabled, + task_config=prompt.task_config, + ) + + async def _render( + self, arguments: dict[str, Any] | None = None + ) -> PromptResult | mcp.types.CreateTaskResult: + """Skip task routing - delegate to render() which calls child middleware. + + The actual underlying prompt will check _task_metadata contextvar and + submit to Docket if appropriate. This wrapper just passes through. + """ + return await self.render(arguments) + + async def render( + self, arguments: dict[str, Any] | None = None + ) -> PromptResult | mcp.types.CreateTaskResult: # type: ignore[override] + """Delegate to child server's middleware. + + When called from a Docket worker (background task), there's no FastMCP + context set up, so we create one for the child server. + + Note: The _docket_fn_key contextvar is intentionally NOT updated here. + The parent set it to the full namespaced name (e.g., c_gc_greet) which + is what the function is registered under in Docket. All provider layers + pass this through unchanged so the eventual prompt._render() uses the + correct Docket lookup key. + """ + import fastmcp.server.context + + try: + from fastmcp.server.dependencies import get_context + + get_context() # Will raise if no context + result = await self._server._get_prompt_content_middleware( + self._original_name, arguments + ) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result + except RuntimeError: + # No context (e.g., Docket worker) - create one for the child server + async with fastmcp.server.context.Context(fastmcp=self._server): + result = await self._server._get_prompt_content_middleware( + self._original_name, arguments + ) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result + + +class FastMCPProviderResourceTemplate(ResourceTemplate): + """Resource template that creates FastMCPProviderResources. + + When `create_resource()` is called, this template creates a + FastMCPProviderResource that will invoke the wrapped server's middleware + when read. + """ + + _server: Any = None # FastMCP, but Any to avoid circular import + _original_uri_template: str | None = None + + def __init__( + self, + server: Any, + original_uri_template: str, + **kwargs: Any, + ): + super().__init__(**kwargs) + self._server = server + self._original_uri_template = original_uri_template + + @classmethod + def wrap( + cls, server: Any, template: ResourceTemplate + ) -> FastMCPProviderResourceTemplate: + """Wrap a ResourceTemplate to create FastMCPProviderResources.""" + return cls( + server=server, + original_uri_template=template.uri_template, + uri_template=template.uri_template, + name=template.name, + description=template.description, + mime_type=template.mime_type, + parameters=template.parameters, + tags=template.tags, + annotations=template.annotations, + enabled=template.enabled, + task_config=template.task_config, + ) + + async def create_resource(self, uri: str, params: dict[str, Any]) -> Resource: + """Create a FastMCPProviderResource for the given URI. + + The `uri` is the external/transformed URI (e.g., with namespace prefix). + We use `_original_uri_template` with `params` to construct the internal + URI that the nested server understands. + """ + # Expand the original template with params to get internal URI + original_uri = _expand_uri_template(self._original_uri_template or "", params) + return FastMCPProviderResource( + server=self._server, + original_uri=original_uri, + uri=AnyUrl(uri), + name=self.name, + description=self.description, + mime_type=self.mime_type, + ) + + async def _read( + self, uri: str, params: dict[str, Any] + ) -> ResourceContent | mcp.types.CreateTaskResult: + """Delegate to child server's middleware. + + Skips task routing at this layer - the child's template._read() will + check _task_metadata contextvar and submit to Docket if appropriate. + + Sets _docket_fn_key to self.uri_template (the transformed pattern) so that + when the child template's _read() submits to Docket, it uses the correct + key that matches what was registered via TransformingProvider.get_tasks(). + + Only sets _docket_fn_key if not already set - in nested mounts, the + outermost wrapper sets the key and inner wrappers preserve it. + """ + import fastmcp.server.context + from fastmcp.server.dependencies import _docket_fn_key + + # Expand the original template with params to get internal URI + original_uri = _expand_uri_template(self._original_uri_template or "", params) + + # Set _docket_fn_key to the template pattern, but only if the current + # value isn't already a template pattern (contains '{'). + # - Server sets concrete URI (e.g., "item://c/gc/42") - no '{', override it + # - Outer wrapper sets pattern (e.g., "item://c/gc/{id}") - has '{', keep it + # In nested mounts (parent→child→grandchild), the outermost wrapper + # has the fully-transformed pattern that matches Docket registration. + existing_key = _docket_fn_key.get() + key_token = None + if not existing_key or "{" not in existing_key: + key_token = _docket_fn_key.set(self.uri_template) + try: + try: + from fastmcp.server.dependencies import get_context + + get_context() # Will raise if no context + result = await self._server._read_resource_middleware(original_uri) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result[0] + except RuntimeError: + # No context (e.g., Docket worker) - create one for the child server + async with fastmcp.server.context.Context(fastmcp=self._server): + result = await self._server._read_resource_middleware(original_uri) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result[0] + finally: + if key_token is not None: + _docket_fn_key.reset(key_token) + + async def read(self, arguments: dict[str, Any]) -> str | bytes: + """Read the resource content for background task execution. + + Creates a resource from this template and reads its content. + This method is called by Docket during background task execution. + """ + # Expand the original template with arguments to get internal URI + original_uri = _expand_uri_template( + self._original_uri_template or "", arguments + ) + + # Create and read the resource + resource = FastMCPProviderResource( + server=self._server, + original_uri=original_uri, + uri=AnyUrl(original_uri), + name=self.name, + description=self.description, + mime_type=self.mime_type, + ) + result = await resource.read() + + # Return raw content (str or bytes) + if hasattr(result, "content"): + return result.content # type: ignore[return-value] + return result # type: ignore[return-value] + + def register_with_docket(self, docket: Docket) -> None: + """No-op: the child's actual template is registered via get_tasks().""" + + async def add_to_docket( # type: ignore[override] + self, + docket: Docket, + params: dict[str, Any], + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, + ) -> Execution: + """Schedule this template for background execution via docket. + + The child's FunctionResourceTemplate.fn is registered (via get_tasks), + and it expects splatted **kwargs, so we splat params here. + """ + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(**params) + + +# ----------------------------------------------------------------------------- +# FastMCPProvider +# ----------------------------------------------------------------------------- class FastMCPProvider(Provider): @@ -32,8 +436,9 @@ class FastMCPProvider(Provider): the mounted server's tools, resources, and prompts through the parent server. - Execution methods (`call_tool`, `read_resource`, `render_prompt`) invoke - the mounted server's middleware chain. + Components returned by this provider are wrapped in FastMCPProvider* + classes that delegate execution to the wrapped server's middleware chain. + This ensures middleware runs when components are executed. Example: ```python @@ -73,48 +478,53 @@ class FastMCPProvider(Provider): # ------------------------------------------------------------------------- async def list_tools(self) -> Sequence[Tool]: - """List all tools from the mounted server.""" - return await self.server._list_tools_middleware() + """List all tools from the mounted server as FastMCPProviderTools. + + Calls the nested server's middleware to list tools, then wraps + each tool as a FastMCPProviderTool that delegates execution to the + nested server's middleware. + """ + raw_tools = await self.server._list_tools_middleware() + return [FastMCPProviderTool.wrap(self.server, t) for t in raw_tools] async def get_tool(self, name: str) -> Tool | None: - """Get a tool by name.""" + """Get a tool by name as a FastMCPProviderTool.""" tools = await self.list_tools() return next((t for t in tools if t.name == name), None) - async def call_tool( - self, name: str, arguments: dict[str, Any] - ) -> ToolResult | None: - """Execute a tool through the mounted server's middleware chain.""" - return await self.server._call_tool_middleware(name, arguments) - # ------------------------------------------------------------------------- # Resource methods # ------------------------------------------------------------------------- async def list_resources(self) -> Sequence[Resource]: - """List all resources from the mounted server.""" - return await self.server._list_resources_middleware() + """List all resources from the mounted server as FastMCPProviderResources. + + Calls the nested server's middleware to list resources, then wraps + each resource as a FastMCPProviderResource that delegates reading to the + nested server's middleware. + """ + raw_resources = await self.server._list_resources_middleware() + return [FastMCPProviderResource.wrap(self.server, r) for r in raw_resources] async def get_resource(self, uri: str) -> Resource | None: - """Get a concrete resource by URI.""" + """Get a concrete resource by URI as a FastMCPProviderResource.""" resources = await self.list_resources() return next((r for r in resources if str(r.uri) == uri), None) - async def read_resource(self, uri: str) -> ResourceContent | None: - """Read a resource through the mounted server's middleware chain.""" - try: - contents = await self.server._read_resource_middleware(uri) - return contents[0] if contents else None - except NotFoundError: - return None - # ------------------------------------------------------------------------- # Resource template methods # ------------------------------------------------------------------------- async def list_resource_templates(self) -> Sequence[ResourceTemplate]: - """List all resource templates from the mounted server.""" - return await self.server._list_resource_templates_middleware() + """List all resource templates from the mounted server. + + Returns FastMCPProviderResourceTemplate instances that create + FastMCPProviderResources when materialized. + """ + raw_templates = await self.server._list_resource_templates_middleware() + return [ + FastMCPProviderResourceTemplate.wrap(self.server, t) for t in raw_templates + ] async def get_resource_template(self, uri: str) -> ResourceTemplate | None: """Get a resource template that matches the given URI.""" @@ -124,30 +534,24 @@ class FastMCPProvider(Provider): return template return None - async def read_resource_template(self, uri: str) -> ResourceContent | None: - """Read a resource via a matching template through the mounted server.""" - # The server's middleware handles template resolution - return await self.read_resource(uri) - # ------------------------------------------------------------------------- # Prompt methods # ------------------------------------------------------------------------- async def list_prompts(self) -> Sequence[Prompt]: - """List all prompts from the mounted server.""" - return await self.server._list_prompts_middleware() + """List all prompts from the mounted server as FastMCPProviderPrompts. + + Returns FastMCPProviderPrompt instances that delegate rendering to the + wrapped server's middleware. + """ + raw_prompts = await self.server._list_prompts_middleware() + return [FastMCPProviderPrompt.wrap(self.server, p) for p in raw_prompts] async def get_prompt(self, name: str) -> Prompt | None: - """Get a prompt by name.""" + """Get a prompt by name as a FastMCPProviderPrompt.""" prompts = await self.list_prompts() return next((p for p in prompts if p.name == name), None) - async def render_prompt( - self, name: str, arguments: dict[str, Any] | None - ) -> PromptResult | None: - """Render a prompt through the mounted server's middleware chain.""" - return await self.server._get_prompt_content_middleware(name, arguments) - # ------------------------------------------------------------------------- # Task registration # ------------------------------------------------------------------------- @@ -155,45 +559,34 @@ class FastMCPProvider(Provider): async def get_tasks(self) -> TaskComponents: """Return task-eligible components from the mounted server. - Accesses the wrapped server's managers directly to avoid triggering - middleware during registration. Also recursively collects tasks from - nested providers. + Returns the child's ACTUAL components (not wrapped) so their actual + functions get registered with Docket. TransformingProvider.get_tasks() + handles namespace transformation of keys. + + Accesses managers directly to avoid triggering middleware during startup. """ - from fastmcp.prompts.prompt import FunctionPrompt - from fastmcp.resources.resource import FunctionResource - from fastmcp.resources.template import FunctionResourceTemplate - from fastmcp.tools.tool import FunctionTool - - tools: list[FunctionTool] = [] - resources: list[FunctionResource] = [] - templates: list[FunctionResourceTemplate] = [] - prompts: list[FunctionPrompt] = [] - - # Direct manager access (bypasses middleware) - for tool in self.server._tool_manager._tools.values(): - if isinstance(tool, FunctionTool) and tool.task_config.mode != "forbidden": - tools.append(tool) - - for resource in self.server._resource_manager._resources.values(): - if ( - isinstance(resource, FunctionResource) - and resource.task_config.mode != "forbidden" - ): - resources.append(resource) - - for template in self.server._resource_manager._templates.values(): - if ( - isinstance(template, FunctionResourceTemplate) - and template.task_config.mode != "forbidden" - ): - templates.append(template) - - for prompt in self.server._prompt_manager._prompts.values(): - if ( - isinstance(prompt, FunctionPrompt) - and prompt.task_config.mode != "forbidden" - ): - prompts.append(prompt) + # Return child's actual components - their .fn gets registered with Docket + # TransformingProvider.get_tasks() transforms keys to include namespace + tools: list[Tool] = [ + t + for t in self.server._tool_manager._tools.values() + if t.task_config.mode != "forbidden" + ] + resources: list[Resource] = [ + r + for r in self.server._resource_manager._resources.values() + if r.task_config.mode != "forbidden" + ] + templates: list[ResourceTemplate] = [ + t + for t in self.server._resource_manager._templates.values() + if t.task_config.mode != "forbidden" + ] + prompts: list[Prompt] = [ + p + for p in self.server._prompt_manager._prompts.values() + if p.task_config.mode != "forbidden" + ] # Recursively get tasks from nested providers for provider in self.server._providers: diff --git a/src/fastmcp/server/proxy.py b/src/fastmcp/server/proxy.py index 2a4ae1136..bd4a1a51e 100644 --- a/src/fastmcp/server/proxy.py +++ b/src/fastmcp/server/proxy.py @@ -185,7 +185,9 @@ class ProxyResourceManager(ResourceManager, ProxyManagerMixin): templates_dict = await self.get_resource_templates() return list(templates_dict.values()) - async def read_resource(self, uri: AnyUrl | str) -> ResourceContent: + async def read_resource( + self, uri: AnyUrl | str + ) -> ResourceContent | mcp.types.CreateTaskResult: """Reads a resource, trying local/mounted first, then proxy if not found.""" try: # First try local and mounted resources @@ -256,7 +258,7 @@ class ProxyPromptManager(PromptManager, ProxyManagerMixin): self, name: str, arguments: dict[str, Any] | None = None, - ) -> PromptResult: + ) -> PromptResult | mcp.types.CreateTaskResult: """Renders a prompt, trying local/mounted first, then proxy if not found.""" try: # First try local and mounted prompts diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index e42bcab68..3ebe56d2e 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -35,12 +35,10 @@ from mcp.server.lowlevel.server import LifespanResultT, NotificationOptions from mcp.server.stdio import stdio_server from mcp.shared.exceptions import McpError from mcp.types import ( - METHOD_NOT_FOUND, Annotations, AnyFunction, CallToolRequestParams, ContentBlock, - ErrorData, GetPromptResult, ToolAnnotations, ) @@ -716,17 +714,6 @@ class FastMCP(Generic[LifespanResultT]): raise NotFoundError(f"Unknown tool: {key}") - async def _get_tool_with_task_config(self, key: str) -> Tool | None: - """Get a tool by key, returning None if not found. - - Used for task config checking where we need the actual tool object - (including from mounted servers and proxies) but don't want to raise. - """ - try: - return await self.get_tool(key) - except NotFoundError: - return None - async def _get_resource_or_template_or_none( self, uri: str ) -> Resource | ResourceTemplate | None: @@ -1266,16 +1253,17 @@ class FastMCP(Generic[LifespanResultT]): """ Handle MCP 'callTool' requests. - Detects SEP-1686 task metadata and routes to background execution if supported. + Sets task metadata contextvar and runs middleware. The tool's _run() method + handles the backgrounding decision, ensuring middleware runs before Docket. Args: key: The name of the tool to call arguments: Arguments to pass to the tool Returns: - List of MCP Content objects containing the tool results + Tool result or CreateTaskResult for background execution """ - from fastmcp.server.tasks.handlers import handle_tool_as_task + from fastmcp.server.dependencies import _docket_fn_key, _task_metadata logger.debug( f"[{self.name}] Handler called: call_tool %s with %s", key, arguments @@ -1283,64 +1271,31 @@ class FastMCP(Generic[LifespanResultT]): async with fastmcp.server.context.Context(fastmcp=self): try: - # Check for SEP-1686 task metadata via request context - task_meta = None + # Extract SEP-1686 task metadata from request context + task_meta_dict: dict[str, Any] | None = None try: - # Access task metadata from SDK's request context ctx = self._mcp_server.request_context if ctx.experimental.is_task: task_meta = ctx.experimental.task_metadata + task_meta_dict = task_meta.model_dump(exclude_none=True) except (AttributeError, LookupError): - # No request context available - proceed without task metadata pass - # Get tool from local manager, mounted servers, or proxy - tool = await self._get_tool_with_task_config(key) - if ( - tool - and self._should_enable_component(tool) - and hasattr(tool, "task_config") - ): - task_mode = tool.task_config.mode # type: ignore[union-attr] + # Set contextvars so tool._run() can access them + task_token = _task_metadata.set(task_meta_dict) + key_token = _docket_fn_key.set(key) + try: + # Middleware always runs - tool._run() handles backgrounding + result = await self._call_tool_middleware(key, arguments) - # Enforce mode="required" - must have task metadata - if task_mode == "required" and not task_meta: - raise McpError( - ErrorData( - code=METHOD_NOT_FOUND, - message=f"Tool '{key}' requires task-augmented execution", - ) - ) + # Result could be CreateTaskResult (from nested tool._run()) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return result.to_mcp_result() + finally: + _task_metadata.reset(task_token) + _docket_fn_key.reset(key_token) - # Route to background if task metadata present and mode allows - if task_meta and task_mode != "forbidden": - # Tool has task support, use Docket for background execution - task_meta_dict = task_meta.model_dump(exclude_none=True) - return await handle_tool_as_task( - self, key, arguments, task_meta_dict - ) - - # Forbidden mode: task requested but mode="forbidden" - # Return error result with returned_immediately=True - if task_meta and task_mode == "forbidden": - return mcp.types.CallToolResult( - content=[ - mcp.types.TextContent( - type="text", - text=f"Tool '{key}' does not support task-augmented execution", - ) - ], - isError=True, - _meta={ - "modelcontextprotocol.io/task": { - "returned_immediately": True - } - }, - ) - - # Synchronous execution (normal path) - result = await self._call_tool_middleware(key, arguments) - return result.to_mcp_result() except DisabledError as e: raise NotFoundError(f"Unknown tool: {key}") from e except NotFoundError as e: @@ -1354,61 +1309,48 @@ class FastMCP(Generic[LifespanResultT]): This is a custom handler because the SDK's read_resource decorator does not support returning CreateTaskResult for background tasks. """ - from fastmcp.server.tasks.handlers import handle_resource_as_task + from fastmcp.server.dependencies import _docket_fn_key, _task_metadata uri = req.params.uri # Check for task metadata via SDK's request context - task_meta = None + task_meta_dict: dict[str, Any] | None = None try: ctx = self._mcp_server.request_context if ctx.experimental.is_task: task_meta = ctx.experimental.task_metadata + task_meta_dict = task_meta.model_dump(exclude_none=True) except (AttributeError, LookupError): pass async with fastmcp.server.context.Context(fastmcp=self): - # Get resource including from mounted servers - resource = await self._get_resource_or_template_or_none(str(uri)) - if ( - resource - and self._should_enable_component(resource) - and hasattr(resource, "task_config") - ): - task_mode = resource.task_config.mode # type: ignore[union-attr] + try: + # Set contextvars so Resource._read() can access them + task_token = _task_metadata.set(task_meta_dict) + key_token = _docket_fn_key.set(str(uri)) + try: + # Middleware always runs - Resource._read() handles backgrounding + result = await self._read_resource_middleware(uri) - # Enforce mode="required" - must have task metadata - if task_mode == "required" and not task_meta: - raise McpError( - ErrorData( - code=METHOD_NOT_FOUND, - message=f"Resource '{uri}' requires task-augmented execution", - ) + # Result could be CreateTaskResult (from nested Resource._read()) + if isinstance(result, mcp.types.CreateTaskResult): + return mcp.types.ServerResult(result) + + # Normal synchronous result + mcp_contents = [ + item.to_mcp_resource_contents(uri) for item in result + ] + return mcp.types.ServerResult( + mcp.types.ReadResourceResult(contents=mcp_contents) ) + finally: + _task_metadata.reset(task_token) + _docket_fn_key.reset(key_token) - # Route to background if task metadata present and mode allows - if task_meta and task_mode != "forbidden": - task_meta_dict = task_meta.model_dump(exclude_none=True) - result = await handle_resource_as_task( - self, str(uri), resource, task_meta_dict - ) - return mcp.types.ServerResult(result) - - # Forbidden mode: task requested but mode="forbidden" - if task_meta and task_mode == "forbidden": - raise McpError( - ErrorData( - code=METHOD_NOT_FOUND, - message=f"Resource '{uri}' does not support task-augmented execution", - ) - ) - - # Synchronous execution - contents = await self._read_resource_mcp(uri) - mcp_contents = [item.to_mcp_resource_contents(uri) for item in contents] - return mcp.types.ServerResult( - mcp.types.ReadResourceResult(contents=mcp_contents) - ) + except DisabledError as e: + raise NotFoundError(f"Unknown resource: {str(uri)!r}") from e + except NotFoundError as e: + raise NotFoundError(f"Unknown resource: {str(uri)!r}") from e async def _get_prompt_handler( self, req: mcp.types.GetPromptRequest @@ -1418,70 +1360,55 @@ class FastMCP(Generic[LifespanResultT]): This is a custom handler because the SDK's get_prompt decorator does not support returning CreateTaskResult for background tasks. """ - from fastmcp.server.tasks.handlers import handle_prompt_as_task + from fastmcp.server.dependencies import _docket_fn_key, _task_metadata name = req.params.name arguments = req.params.arguments # Check for task metadata via SDK's request context - task_meta = None + task_meta_dict: dict[str, Any] | None = None try: ctx = self._mcp_server.request_context if ctx.experimental.is_task: task_meta = ctx.experimental.task_metadata + task_meta_dict = task_meta.model_dump(exclude_none=True) except (AttributeError, LookupError): pass async with fastmcp.server.context.Context(fastmcp=self): try: - prompt = await self.get_prompt(name) - except NotFoundError: - prompt = None - if ( - prompt - and self._should_enable_component(prompt) - and hasattr(prompt, "task_config") - and prompt.task_config - ): - task_mode = prompt.task_config.mode # type: ignore[union-attr] + # Set contextvars so Prompt._render() can access them + task_token = _task_metadata.set(task_meta_dict) + key_token = _docket_fn_key.set(name) + try: + # Middleware always runs - Prompt._render() handles backgrounding + result = await self._get_prompt_content_middleware(name, arguments) - # Enforce mode="required" - must have task metadata - if task_mode == "required" and not task_meta: - raise McpError( - ErrorData( - code=METHOD_NOT_FOUND, - message=f"Prompt '{name}' requires task-augmented execution", - ) - ) + # Result could be CreateTaskResult (from nested Prompt._render()) + if isinstance(result, mcp.types.CreateTaskResult): + return mcp.types.ServerResult(result) - # Route to background if task metadata present and mode allows - if task_meta and task_mode != "forbidden": - task_meta_dict = task_meta.model_dump(exclude_none=True) - result = await handle_prompt_as_task( - self, name, arguments, task_meta_dict - ) - return mcp.types.ServerResult(result) + # Normal synchronous result + return mcp.types.ServerResult(result.to_mcp_prompt_result()) + finally: + _task_metadata.reset(task_token) + _docket_fn_key.reset(key_token) - # Forbidden mode: task requested but mode="forbidden" - if task_meta and task_mode == "forbidden": - raise McpError( - ErrorData( - code=METHOD_NOT_FOUND, - message=f"Prompt '{name}' does not support task-augmented execution", - ) - ) - - # Synchronous execution - result = await self._get_prompt_mcp(name, arguments) - return mcp.types.ServerResult(result) + except DisabledError as e: + raise NotFoundError(f"Unknown prompt: {name!r}") from e + except NotFoundError as e: + raise NotFoundError(f"Unknown prompt: {name!r}") from e async def _call_tool_middleware( self, key: str, arguments: dict[str, Any], - ) -> ToolResult: + ) -> ToolResult | mcp.types.CreateTaskResult: """ Applies this server's middleware and delegates the filtered call to the manager. + + Returns ToolResult for synchronous execution, or CreateTaskResult if the + tool was submitted to Docket for background execution. """ mw_context = MiddlewareContext[CallToolRequestParams]( @@ -1498,7 +1425,7 @@ class FastMCP(Generic[LifespanResultT]): async def _call_tool( self, context: MiddlewareContext[mcp.types.CallToolRequestParams], - ) -> ToolResult: + ) -> ToolResult | mcp.types.CreateTaskResult: """ Call a tool """ @@ -1516,35 +1443,24 @@ class FastMCP(Generic[LifespanResultT]): # Try component providers (first registered wins) for provider in self._providers: - try: - tool = await provider.get_tool(tool_name) - if tool is not None and self._should_enable_component(tool): - result = await provider.call_tool( - tool_name, context.message.arguments or {} - ) - if result is not None: - return result - except (ValidationError, PydanticValidationError): - # Validation errors are never masked - logger.exception(f"Error validating tool {tool_name!r}") - raise - except ToolError: - logger.exception(f"Error calling tool {tool_name!r}") - raise - except Exception as e: - logger.exception(f"Error calling tool {tool_name!r} from provider") - if self._mask_error_details: - raise ToolError(f"Error calling tool {tool_name!r}") from e - raise ToolError(f"Error calling tool {tool_name!r}: {e}") from e + tool = await provider.get_tool(tool_name) + if tool is not None and self._should_enable_component(tool): + return await self._execute_tool( + tool, tool_name, context.message.arguments or {} + ) raise NotFoundError(f"Unknown tool: {tool_name!r}") async def _execute_tool( self, tool: Tool, tool_name: str, arguments: dict[str, Any] - ) -> ToolResult: - """Run a tool with unified error handling.""" + ) -> ToolResult | mcp.types.CreateTaskResult: + """Run a tool with unified error handling. + + Calls tool._run() which handles task routing - checking the task_metadata + contextvar and submitting to Docket if appropriate. + """ try: - return await tool.run(arguments) + return await tool._run(arguments) except (ValidationError, PydanticValidationError): # Validation errors are never masked - they indicate client input issues logger.exception(f"Error validating tool {tool_name!r}") @@ -1569,7 +1485,14 @@ class FastMCP(Generic[LifespanResultT]): async with fastmcp.server.context.Context(fastmcp=self): try: # Task routing handled by custom handler - return list[ResourceContent](await self._read_resource_middleware(uri)) + # Note: Without task metadata, _read_resource_middleware always returns list + result = await self._read_resource_middleware(uri) + if isinstance(result, mcp.types.CreateTaskResult): + # Should never happen without task metadata, but handle for type safety + raise RuntimeError( + "Unexpected CreateTaskResult in _read_resource_mcp" + ) + return result except DisabledError as e: # convert to NotFoundError to avoid leaking resource presence raise NotFoundError(f"Unknown resource: {str(uri)!r}") from e @@ -1580,9 +1503,12 @@ class FastMCP(Generic[LifespanResultT]): async def _read_resource_middleware( self, uri: AnyUrl | str, - ) -> list[ResourceContent]: + ) -> list[ResourceContent] | mcp.types.CreateTaskResult: """ Applies this server's middleware and delegates the filtered call to the manager. + + Returns list[ResourceContent] for synchronous execution, or CreateTaskResult + if the resource was submitted to Docket for background execution. """ # Convert string URI to AnyUrl if needed @@ -1595,32 +1521,41 @@ class FastMCP(Generic[LifespanResultT]): method="resources/read", fastmcp_context=fastmcp.server.dependencies.get_context(), ) - return list( - await self._apply_middleware( - context=mw_context, call_next=self._read_resource - ) + result = await self._apply_middleware( + context=mw_context, call_next=self._read_resource ) + # CreateTaskResult passes through, otherwise convert to list + if isinstance(result, mcp.types.CreateTaskResult): + return result + return list(result) async def _read_resource( self, context: MiddlewareContext[mcp.types.ReadResourceRequestParams], - ) -> list[ResourceContent]: + ) -> list[ResourceContent] | mcp.types.CreateTaskResult: """ - Read a resource + Read a resource. + + Returns list[ResourceContent] for synchronous execution, or CreateTaskResult + if the resource was submitted to Docket for background execution. """ + uri_str = str(context.message.uri) - # Try local resources first (static resources take precedence) - try: - resource = await self._resource_manager.get_resource(uri_str) + # Try local concrete resources first (static resources take precedence) + # Note: Don't use get_resource() here because it creates resources from templates, + # which would bypass our template execution flow that handles task routing properly. + local_resources = await self._resource_manager.get_resources() + if uri_str in local_resources: + resource = local_resources[uri_str] if self._should_enable_component(resource): - content = await self._execute_resource(resource, uri_str) + result = await self._execute_resource(resource, uri_str) + if isinstance(result, mcp.types.CreateTaskResult): + return result # Use mime_type from ResourceContent if set, otherwise from resource - if content.mime_type is None: - content.mime_type = resource.mime_type - return [content] - except NotFoundError: - pass + if result.mime_type is None: + result.mime_type = resource.mime_type + return [result] # Try local templates templates = await self._resource_manager.get_resource_templates() @@ -1628,55 +1563,72 @@ class FastMCP(Generic[LifespanResultT]): params = template.matches(uri_str) if params is not None: if self._should_enable_component(template): - resource = await template.create_resource(uri_str, params) - content = await self._execute_resource(resource, uri_str) - return [content] + # Templates need special task routing - call _execute_template + # which handles passing params to Docket + result = await self._execute_template(template, uri_str, params) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return [result] # Try component providers (first registered wins) - concrete resources for provider in self._providers: - try: - resource = await provider.get_resource(uri_str) - if resource is not None and self._should_enable_component(resource): - content = await provider.read_resource(uri_str) - if content is not None: - return [content] - except ResourceError: - logger.exception(f"Error reading resource {uri_str!r}") - raise - except Exception as e: - logger.exception(f"Error reading resource {uri_str!r} from provider") - if self._mask_error_details: - raise ResourceError(f"Error reading resource {uri_str!r}") from e - raise ResourceError(f"Error reading resource {uri_str!r}: {e}") from e + resource = await provider.get_resource(uri_str) + if resource is not None and self._should_enable_component(resource): + result = await self._execute_resource(resource, uri_str) + if isinstance(result, mcp.types.CreateTaskResult): + return result + if result.mime_type is None: + result.mime_type = resource.mime_type + return [result] # Try component providers (first registered wins) - templates for provider in self._providers: - try: - template = await provider.get_resource_template(uri_str) - if template is not None and self._should_enable_component(template): - content = await provider.read_resource_template(uri_str) - if content is not None: - return [content] - except ResourceError: - logger.exception(f"Error reading resource {uri_str!r}") - raise - except Exception as e: - logger.exception( - f"Error reading resource {uri_str!r} from provider template" - ) - if self._mask_error_details: - raise ResourceError(f"Error reading resource {uri_str!r}") from e - raise ResourceError(f"Error reading resource {uri_str!r}: {e}") from e + template = await provider.get_resource_template(uri_str) + if template is not None and self._should_enable_component(template): + params = template.matches(uri_str) + if params is not None: + # Templates need special task routing - call _execute_template + # which handles passing params to Docket + result = await self._execute_template(template, uri_str, params) + if isinstance(result, mcp.types.CreateTaskResult): + return result + return [result] raise NotFoundError(f"Unknown resource: {uri_str!r}") async def _execute_resource( self, resource: Resource, uri_str: str - ) -> ResourceContent: - """Read a resource with unified error handling.""" + ) -> ResourceContent | mcp.types.CreateTaskResult: + """Read a resource with unified error handling. + + Calls resource._read() which handles task routing - checking the task_metadata + contextvar and submitting to Docket if appropriate. + """ try: return await resource._read() - except ResourceError: + except (ResourceError, McpError): + logger.exception(f"Error reading resource {uri_str!r}") + raise + except Exception as e: + logger.exception(f"Error reading resource {uri_str!r}") + if self._mask_error_details: + raise ResourceError(f"Error reading resource {uri_str!r}") from e + raise ResourceError(f"Error reading resource {uri_str!r}: {e}") from e + + async def _execute_template( + self, + template: ResourceTemplate, + uri_str: str, + params: dict[str, Any], + ) -> ResourceContent | mcp.types.CreateTaskResult: + """Execute a template with unified error handling. + + Calls template._read() which handles task routing - checking the task_metadata + contextvar and submitting to Docket if appropriate. + """ + try: + return await template._read(uri_str, params) + except (ResourceError, McpError): logger.exception(f"Error reading resource {uri_str!r}") raise except Exception as e: @@ -1716,16 +1668,27 @@ class FastMCP(Generic[LifespanResultT]): """ Applies this server's middleware and delegates the filtered call to the manager. Converts PromptResult to GetPromptResult for MCP protocol. + + Note: This method assumes synchronous execution. For task-augmented execution, + use _get_prompt_content_middleware directly and handle CreateTaskResult. """ result = await self._get_prompt_content_middleware(name, arguments) + if isinstance(result, mcp.types.CreateTaskResult): + raise RuntimeError( + "Prompt returned CreateTaskResult but _get_prompt_middleware " + "expects synchronous execution" + ) return result.to_mcp_prompt_result() async def _get_prompt_content_middleware( self, name: str, arguments: dict[str, Any] | None = None - ) -> PromptResult: + ) -> PromptResult | mcp.types.CreateTaskResult: """ Applies this server's middleware and returns PromptResult. Used internally and by parent servers for mounted prompts. + + Returns PromptResult for synchronous execution, or CreateTaskResult + if the prompt was submitted to Docket for background execution. """ mw_context = MiddlewareContext( message=mcp.types.GetPromptRequestParams(name=name, arguments=arguments), @@ -1741,7 +1704,13 @@ class FastMCP(Generic[LifespanResultT]): async def _get_prompt( self, context: MiddlewareContext[mcp.types.GetPromptRequestParams], - ) -> PromptResult: + ) -> PromptResult | mcp.types.CreateTaskResult: + """ + Get a prompt. + + Returns PromptResult for synchronous execution, or CreateTaskResult + if the prompt was submitted to Docket for background execution. + """ name = context.message.name # Try local prompts first (static prompts take precedence) @@ -1756,32 +1725,25 @@ class FastMCP(Generic[LifespanResultT]): # Try component providers (first registered wins) for provider in self._providers: - try: - prompt = await provider.get_prompt(name) - if prompt is not None and self._should_enable_component(prompt): - result = await provider.render_prompt( - name, context.message.arguments - ) - if result is not None: - return result - except PromptError: - logger.exception(f"Error rendering prompt {name!r}") - raise - except Exception as e: - logger.exception(f"Error rendering prompt {name!r} from provider") - if self._mask_error_details: - raise PromptError(f"Error rendering prompt {name!r}") from e - raise PromptError(f"Error rendering prompt {name!r}: {e}") from e + prompt = await provider.get_prompt(name) + if prompt is not None and self._should_enable_component(prompt): + return await self._execute_prompt( + prompt, name, context.message.arguments + ) raise NotFoundError(f"Unknown prompt: {name!r}") async def _execute_prompt( self, prompt: Prompt, name: str, arguments: dict[str, Any] | None - ) -> PromptResult: - """Render a prompt with unified error handling.""" + ) -> PromptResult | mcp.types.CreateTaskResult: + """Render a prompt with unified error handling. + + Calls prompt._render() which handles task routing - checking the task_metadata + contextvar and submitting to Docket if appropriate. + """ try: return await prompt._render(arguments) - except PromptError: + except (PromptError, McpError): logger.exception(f"Error rendering prompt {name!r}") raise except Exception as e: diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index dd33d1374..6c0b424e5 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/src/fastmcp/server/tasks/handlers.py @@ -8,7 +8,7 @@ from __future__ import annotations import uuid from contextlib import suppress from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import mcp.types from mcp.shared.exceptions import McpError @@ -18,34 +18,38 @@ from fastmcp.server.dependencies import _current_docket, get_context from fastmcp.server.tasks.keys import build_task_key if TYPE_CHECKING: + from fastmcp.prompts.prompt import Prompt from fastmcp.resources.resource import Resource from fastmcp.resources.template import ResourceTemplate - from fastmcp.server.server import FastMCP + from fastmcp.tools.tool import Tool # Redis mapping TTL buffer: Add 15 minutes to Docket's execution_ttl TASK_MAPPING_TTL_BUFFER_SECONDS = 15 * 60 -async def handle_tool_as_task( - server: FastMCP, - tool_name: str, - arguments: dict[str, Any], - _task_meta: dict[str, Any], +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, ) -> mcp.types.CreateTaskResult: - """Handle tool execution as background task (SEP-1686). + """Submit any component to Docket for background execution (SEP-1686). - Queues the user's actual function to Docket (preserving signature for DI), - stores raw return values, converts to MCP types on retrieval. + 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. Note: Client-requested TTL in task_meta is intentionally ignored. Server-side TTL policy (docket.execution_ttl) takes precedence for consistent task lifecycle management. Args: - server: FastMCP server instance - tool_name: Name of the tool to execute - arguments: Tool arguments - _task_meta: Task metadata from request (unused - server TTL policy applies) + 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) Returns: CreateTaskResult: Task stub with proper Task object @@ -71,10 +75,7 @@ async def handle_tool_as_task( ) # Build full task key with embedded metadata - task_key = build_task_key(session_id, server_task_id, "tool", tool_name) - - # Get the tool to access user's function - tool = await server.get_tool(tool_name) + task_key = build_task_key(session_id, server_task_id, task_type, key) # Store task key mapping and creation timestamp in Redis for protocol handlers redis_key = f"fastmcp:task:{session_id}:{server_task_id}" @@ -98,230 +99,19 @@ async def handle_tool_as_task( } }, ) - - ctx = get_context() with suppress(Exception): # Don't let notification failures break task creation await ctx.session.send_notification(notification) # type: ignore[arg-type] # Queue function to Docket by key (result storage via execution_ttl) - # Use tool.add_to_docket() which handles calling conventions - await tool.add_to_docket(docket, arguments, key=task_key) - - # Spawn subscription task to send status notifications (SEP-1686 optional feature) - from fastmcp.server.tasks.subscriptions import subscribe_to_task_updates - - # Start subscription in session's task group (persists for connection lifetime) - if hasattr(ctx.session, "_subscription_task_group"): - tg = ctx.session._subscription_task_group # type: ignore[attr-defined] - if tg: - tg.start_soon( # type: ignore[union-attr] - subscribe_to_task_updates, - server_task_id, - task_key, - ctx.session, - docket, - ) - - # 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( - taskId=server_task_id, - status="working", - createdAt=created_at, - lastUpdatedAt=created_at, - ttl=int(docket.execution_ttl.total_seconds() * 1000), - pollInterval=1000, - ) - ) - - -async def handle_prompt_as_task( - server: FastMCP, - prompt_name: str, - arguments: dict[str, Any] | None, - _task_meta: dict[str, Any], -) -> mcp.types.CreateTaskResult: - """Handle prompt execution as background task (SEP-1686). - - Queues the user's actual function to Docket (preserving signature for DI). - - Note: Client-requested TTL in task_meta is intentionally ignored. - Server-side TTL policy (docket.execution_ttl) takes precedence. - - Args: - server: FastMCP server instance - prompt_name: Name of the prompt to execute - arguments: Prompt arguments - _task_meta: Task metadata from request (unused - server TTL policy applies) - - Returns: - CreateTaskResult: Task stub with proper Task object - """ - # 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) - created_at = datetime.now(timezone.utc) - - # Get session ID and Docket - ctx = get_context() - session_id = ctx.session_id - - docket = _current_docket.get() - if docket is None: - raise McpError( - ErrorData( - code=INTERNAL_ERROR, - message="Background tasks require a running FastMCP server context", - ) - ) - - # Build full task key with embedded metadata - task_key = build_task_key(session_id, server_task_id, "prompt", prompt_name) - - # Get the prompt - prompt = await server.get_prompt(prompt_name) - - # Store task key mapping and creation timestamp in Redis for protocol handlers - redis_key = f"fastmcp:task:{session_id}:{server_task_id}" - created_at_key = f"fastmcp:task:{session_id}:{server_task_id}:created_at" - ttl_seconds = int( - docket.execution_ttl.total_seconds() + TASK_MAPPING_TTL_BUFFER_SECONDS - ) - async with docket.redis() as redis: - await redis.set(redis_key, task_key, ex=ttl_seconds) - await redis.set(created_at_key, created_at.isoformat(), ex=ttl_seconds) - - # Send notifications/tasks/created per SEP-1686 (mandatory) - # Send BEFORE queuing to avoid race where task completes before notification - notification = mcp.types.JSONRPCNotification( - jsonrpc="2.0", - method="notifications/tasks/created", - params={}, - _meta={ - "modelcontextprotocol.io/related-task": { - "taskId": server_task_id, - } - }, - ) - with suppress(Exception): - await ctx.session.send_notification(notification) # type: ignore[arg-type] - - # Queue function to Docket by key (result storage via execution_ttl) - # Use prompt.add_to_docket() which handles calling conventions - await prompt.add_to_docket(docket, arguments, key=task_key) - - # Spawn subscription task to send status notifications (SEP-1686 optional feature) - from fastmcp.server.tasks.subscriptions import subscribe_to_task_updates - - # Start subscription in session's task group (persists for connection lifetime) - if hasattr(ctx.session, "_subscription_task_group"): - tg = ctx.session._subscription_task_group # type: ignore[attr-defined] - if tg: - tg.start_soon( # type: ignore[union-attr] - subscribe_to_task_updates, - server_task_id, - task_key, - ctx.session, - docket, - ) - - # 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( - taskId=server_task_id, - status="working", - createdAt=created_at, - lastUpdatedAt=created_at, - ttl=int(docket.execution_ttl.total_seconds() * 1000), - pollInterval=1000, - ) - ) - - -async def handle_resource_as_task( - _server: FastMCP, - uri: str, - resource: Resource | ResourceTemplate, - _task_meta: dict[str, Any], -) -> mcp.types.CreateTaskResult: - """Handle resource read as background task (SEP-1686). - - Queues the user's actual function to Docket. - - Note: Client-requested TTL in task_meta is intentionally ignored. - Server-side TTL policy (docket.execution_ttl) takes precedence. - - Args: - _server: FastMCP server instance (unused - kept for signature consistency) - uri: Resource URI - resource: Resource or ResourceTemplate object - _task_meta: Task metadata from request (unused - server TTL policy applies) - - Returns: - CreateTaskResult: Task stub with proper Task object - """ - # 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) - created_at = datetime.now(timezone.utc) - - # Get session ID and Docket - ctx = get_context() - session_id = ctx.session_id - - docket = _current_docket.get() - if docket is None: - raise McpError( - ErrorData( - code=INTERNAL_ERROR, - message="Background tasks require Docket", - ) - ) - - # Build full task key with embedded metadata (use original URI) - task_key = build_task_key(session_id, server_task_id, "resource", str(uri)) - - # Store task key mapping and creation timestamp in Redis for protocol handlers - redis_key = f"fastmcp:task:{session_id}:{server_task_id}" - created_at_key = f"fastmcp:task:{session_id}:{server_task_id}:created_at" - ttl_seconds = int( - docket.execution_ttl.total_seconds() + TASK_MAPPING_TTL_BUFFER_SECONDS - ) - async with docket.redis() as redis: - await redis.set(redis_key, task_key, ex=ttl_seconds) - await redis.set(created_at_key, created_at.isoformat(), ex=ttl_seconds) - - # Send notifications/tasks/created per SEP-1686 (mandatory) - # Send BEFORE queuing to avoid race where task completes before notification - notification = mcp.types.JSONRPCNotification( - jsonrpc="2.0", - method="notifications/tasks/created", - params={}, - _meta={ - "modelcontextprotocol.io/related-task": { - "taskId": server_task_id, - } - }, - ) - with suppress(Exception): - await ctx.session.send_notification(notification) # type: ignore[arg-type] - - # Queue function to Docket by key (result storage via execution_ttl) - # Use add_to_docket() which handles calling conventions - from fastmcp.resources.template import ResourceTemplate, match_uri_template - - if isinstance(resource, ResourceTemplate): - params = match_uri_template(uri, resource.uri_template) or {} - await resource.add_to_docket(docket, params, key=task_key) + # 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:{session}:{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 component.add_to_docket(docket, fn_key=key, task_key=task_key) # type: ignore[call-arg] else: - await resource.add_to_docket(docket, key=task_key) + await component.add_to_docket(docket, arguments, fn_key=key, task_key=task_key) # type: ignore[call-arg] # Spawn subscription task to send status notifications (SEP-1686 optional feature) from fastmcp.server.tasks.subscriptions import subscribe_to_task_updates diff --git a/src/fastmcp/server/tasks/protocol.py b/src/fastmcp/server/tasks/protocol.py index bf02a43d6..4e8d6a998 100644 --- a/src/fastmcp/server/tasks/protocol.py +++ b/src/fastmcp/server/tasks/protocol.py @@ -234,7 +234,7 @@ async def tasks_result_handler(server: FastMCP, params: dict[str, Any]) -> Any: } } - # Convert based on task type using component.convert_result() + to_mcp_result() + # Convert based on task type if task_type == "tool": tool = await server.get_tool(component_id) fastmcp_result = tool.convert_result(raw_value) @@ -263,14 +263,17 @@ async def tasks_result_handler(server: FastMCP, params: dict[str, Any]) -> Any: return mcp_result elif task_type == "resource": - # Convert raw value to ResourceContent (handles str, bytes, ResourceContent) - from fastmcp.resources.resource import ResourceContent - - if isinstance(raw_value, ResourceContent): - resource_content = raw_value - else: - resource_content = ResourceContent.from_value(raw_value) + resource = await server.get_resource(component_id) + resource_content = resource.convert_result(raw_value) + mcp_content = resource_content.to_mcp_resource_contents(component_id) + return mcp.types.ReadResourceResult( + contents=[mcp_content], + _meta=related_task_meta, + ) + elif task_type == "template": + template = await server.get_resource_template(component_id) + resource_content = template.convert_result(raw_value) mcp_content = resource_content.to_mcp_resource_contents(component_id) return mcp.types.ReadResourceResult( contents=[mcp_content], diff --git a/src/fastmcp/server/tasks/routing.py b/src/fastmcp/server/tasks/routing.py new file mode 100644 index 000000000..f0e0d804c --- /dev/null +++ b/src/fastmcp/server/tasks/routing.py @@ -0,0 +1,75 @@ +"""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, ErrorData + +from fastmcp.server.dependencies import get_task_metadata +from fastmcp.server.tasks.handlers import submit_to_docket + +if TYPE_CHECKING: + from fastmcp.prompts.prompt import Prompt + from fastmcp.resources.resource import Resource + from fastmcp.resources.template import ResourceTemplate + from fastmcp.tools.tool import Tool + +TaskType = Literal["tool", "resource", "template", "prompt"] + + +async def check_background_task( + component: Tool | Resource | ResourceTemplate | Prompt, + task_type: TaskType, + key: str, + arguments: dict[str, Any] | 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") + key: Docket registration key (caller resolves from contextvar + fallback) + arguments: Arguments for tool/prompt/template execution + + 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_meta = get_task_metadata() + 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( + ErrorData( + code=METHOD_NOT_FOUND, + message=f"{entity_label} requires task-augmented execution", + ) + ) + + # Enforce mode="forbidden" - cannot be called with task metadata + if task_config.mode == "forbidden" and task_meta: + raise McpError( + ErrorData( + 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 + + return await submit_to_docket(task_type, key, component, arguments) diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 4adafe0ae..3a8baed6e 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -241,11 +241,69 @@ class Tool(FastMCPComponent): raise NotImplementedError("Subclasses must implement run()") def convert_result(self, raw_value: Any) -> ToolResult: - """Convert a raw return value to ToolResult. + """Convert a raw result to ToolResult. - Subclasses should override this to handle their specific conversion logic. + Handles ToolResult passthrough and converts raw values using the tool's + attributes (serializer, output_schema) for proper conversion. """ - raise NotImplementedError("Subclasses must implement convert_result()") + if isinstance(raw_value, ToolResult): + return raw_value + + content = _convert_to_content(raw_value, serializer=self.serializer) + + # Skip structured content for ContentBlock types only if no output_schema + # (if output_schema exists, MCP SDK requires structured_content) + if self.output_schema is None and ( + isinstance(raw_value, ContentBlock | Audio | Image | File) + or ( + isinstance(raw_value, list | tuple) + and any(isinstance(item, ContentBlock) for item in raw_value) + ) + ): + return ToolResult(content=content) + + try: + structured = pydantic_core.to_jsonable_python(raw_value) + except pydantic_core.PydanticSerializationError: + return ToolResult(content=content) + + if self.output_schema is None: + # No schema - only use structured_content for dicts + if isinstance(structured, dict): + return ToolResult(content=content, structured_content=structured) + return ToolResult(content=content) + + # Has output_schema - wrap if x-fastmcp-wrap-result is set + wrap_result = self.output_schema.get("x-fastmcp-wrap-result") + return ToolResult( + content=content, + structured_content={"result": structured} if wrap_result else structured, + ) + + async def _run( + self, arguments: dict[str, Any] + ) -> ToolResult | mcp.types.CreateTaskResult: + """Server entry point that handles task routing. + + This allows ANY Tool subclass to support background execution by setting + task_config.mode to "supported" or "required". The server calls this + method instead of run() directly. + + Subclasses can override this to customize task routing behavior. + For example, FastMCPProviderTool overrides to delegate to child + middleware without submitting to Docket. + """ + from fastmcp.server.dependencies import _docket_fn_key + from fastmcp.server.tasks.routing import check_background_task + + key = _docket_fn_key.get() or self.key + task_result = await check_background_task( + component=self, task_type="tool", key=key, arguments=arguments + ) + if task_result: + return task_result + + return await self.run(arguments) def register_with_docket(self, docket: Docket) -> None: """Register this tool with docket for background execution.""" @@ -254,10 +312,27 @@ class Tool(FastMCPComponent): docket.register(self.run, names=[self.key]) async def add_to_docket( # type: ignore[override] - self, docket: Docket, arguments: dict[str, Any], **kwargs: Any + self, + docket: Docket, + arguments: dict[str, Any], + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: - """Schedule this tool for background execution via docket.""" - return await docket.add(self.key, **kwargs)(arguments) + """Schedule this tool for background execution via docket. + + Args: + docket: The Docket instance + arguments: Tool arguments + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() + """ + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(arguments) @classmethod def from_tool( @@ -402,46 +477,6 @@ class FunctionTool(Tool): return self.convert_result(result) - def convert_result(self, raw_value: Any) -> ToolResult: - """Convert a raw return value to ToolResult. - - This handles the same conversion logic as run(), but works on - already-executed raw values (e.g., from Docket background execution). - """ - if isinstance(raw_value, ToolResult): - return raw_value - - unstructured_result = _convert_to_content(raw_value, serializer=self.serializer) - - if self.output_schema is None: - # Do not produce a structured output for MCP Content Types - if isinstance(raw_value, ContentBlock | Audio | Image | File) or ( - isinstance(raw_value, list | tuple) - and any(isinstance(item, ContentBlock) for item in raw_value) - ): - return ToolResult(content=unstructured_result) - - # Otherwise, try to serialize the result as a dict - try: - structured_content = pydantic_core.to_jsonable_python(raw_value) - if isinstance(structured_content, dict): - return ToolResult( - content=unstructured_result, - structured_content=structured_content, - ) - - except pydantic_core.PydanticSerializationError: - pass - - return ToolResult(content=unstructured_result) - - wrap_result = self.output_schema.get("x-fastmcp-wrap-result") - - return ToolResult( - content=unstructured_result, - structured_content={"result": raw_value} if wrap_result else raw_value, - ) - def register_with_docket(self, docket: Docket) -> None: """Register this tool with docket for background execution. @@ -453,13 +488,29 @@ class FunctionTool(Tool): docket.register(self.fn, names=[self.key]) async def add_to_docket( # type: ignore[override] - self, docket: Docket, arguments: dict[str, Any], **kwargs: Any + self, + docket: Docket, + arguments: dict[str, Any], + *, + fn_key: str | None = None, + task_key: str | None = None, + **kwargs: Any, ) -> Execution: """Schedule this tool for background execution via docket. FunctionTool splats the arguments dict since .fn expects **kwargs. + + Args: + docket: The Docket instance + arguments: Tool arguments + fn_key: Function lookup key in Docket registry (defaults to self.key) + task_key: Redis storage key for the result + **kwargs: Additional kwargs passed to docket.add() """ - return await docket.add(self.key, **kwargs)(**arguments) + lookup_key = fn_key or self.key + if task_key: + kwargs["key"] = task_key + return await docket.add(lookup_key, **kwargs)(**arguments) def _is_object_schema(schema: dict[str, Any]) -> bool: diff --git a/tests/contrib/test_component_manager.py b/tests/contrib/test_component_manager.py index 56be8f625..e0fd2f04a 100644 --- a/tests/contrib/test_component_manager.py +++ b/tests/contrib/test_component_manager.py @@ -42,7 +42,7 @@ class TestComponentManagementRoutes: def mcp(self, mounted_mcp): """Create a FastMCP server with test tools, resources, and prompts.""" mcp = FastMCP("TestServer") - mcp.mount(mounted_mcp, prefix="sub") + mcp.mount(mounted_mcp, namespace="sub") set_up_component_manager(server=mcp) # Add a test tool diff --git a/tests/server/middleware/test_middleware.py b/tests/server/middleware/test_middleware.py index d93b6a202..46dda3181 100644 --- a/tests/server/middleware/test_middleware.py +++ b/tests/server/middleware/test_middleware.py @@ -506,7 +506,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.call_tool("add", {"a": 1, "b": 2}) @@ -526,7 +526,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.call_tool("nested_add", {"a": 1, "b": 2}) @@ -550,7 +550,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.read_resource("resource://test") @@ -570,7 +570,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.read_resource("resource://nested/test") @@ -594,7 +594,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.read_resource("resource://test-template/1") @@ -614,7 +614,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.read_resource("resource://nested/test-template/1") @@ -638,7 +638,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.get_prompt("test_prompt", {"x": "test"}) @@ -658,7 +658,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.get_prompt("nested_test_prompt", {"x": "test"}) @@ -682,7 +682,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.list_tools() @@ -706,7 +706,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.list_resources() @@ -730,7 +730,7 @@ class TestNestedMiddlewareHooks: recording_middleware: RecordingMiddleware, nested_middleware: RecordingMiddleware, ): - mcp_server.mount(nested_mcp_server, prefix="nested") + mcp_server.mount(nested_mcp_server, namespace="nested") async with Client(mcp_server) as client: await client.list_resource_templates() diff --git a/tests/server/providers/test_fastmcp_provider.py b/tests/server/providers/test_fastmcp_provider.py index 11b0cb924..0ef7b266b 100644 --- a/tests/server/providers/test_fastmcp_provider.py +++ b/tests/server/providers/test_fastmcp_provider.py @@ -1,8 +1,73 @@ """Tests for FastMCPProvider.""" +from collections.abc import Sequence + +import mcp.types as mt + from fastmcp import FastMCP from fastmcp.client import Client +from fastmcp.prompts.prompt import PromptResult +from fastmcp.resources.resource import ResourceContent +from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.server.providers import FastMCPProvider +from fastmcp.tools.tool import ToolResult + + +class ToolTracingMiddleware(Middleware): + """Middleware that traces tool calls.""" + + def __init__(self, name: str, calls: list[str]): + super().__init__() + self._name = name + self._calls = calls + + async def on_call_tool( + self, + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, ToolResult], + ) -> ToolResult: + self._calls.append(f"{self._name}:before") + result = await call_next(context) + self._calls.append(f"{self._name}:after") + 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, Sequence[ResourceContent]], + ) -> Sequence[ResourceContent]: + 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 TestToolOperations: @@ -231,3 +296,147 @@ class TestServerReference: provider = FastMCPProvider(server) assert provider.server.name == "MyServer" + + +class TestMiddlewareChain: + """Test that middleware runs at each level of mounted servers.""" + + async def test_tool_middleware_three_levels(self): + """Middleware runs at parent, child, and grandchild levels for tools.""" + calls: list[str] = [] + + grandchild = FastMCP("Grandchild") + + @grandchild.tool + async def compute(x: int) -> int: + calls.append("grandchild:tool") + 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.mount(child, namespace="c") + parent.add_middleware(ToolTracingMiddleware("parent", calls)) + + async with Client(parent) as client: + result = await client.call_tool("c_gc_compute", {"x": 5}) + assert result.data == 10 + + assert calls == [ + "parent:before", + "child:before", + "grandchild:before", + "grandchild:tool", + "grandchild:after", + "child:after", + "parent:after", + ] + + async def test_resource_middleware_three_levels(self): + """Middleware runs at parent, child, and grandchild levels for resources.""" + calls: list[str] = [] + + grandchild = FastMCP("Grandchild") + + @grandchild.resource("data://value") + 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) as client: + result = await client.read_resource("data://c/gc/value") + assert result[0].text == "result" # type: ignore[attr-defined] + + assert calls == [ + "parent:before", + "child:before", + "grandchild:before", + "grandchild:resource", + "grandchild:after", + "child:after", + "parent:after", + ] + + async def test_prompt_middleware_three_levels(self): + """Middleware runs at parent, child, and grandchild levels for prompts.""" + calls: list[str] = [] + + grandchild = FastMCP("Grandchild") + + @grandchild.prompt + 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) as client: + result = await client.get_prompt("c_gc_greet", {"name": "World"}) + assert result.messages[0].content.text == "Hello, World!" # type: ignore[attr-defined] + + assert calls == [ + "parent:before", + "child:before", + "grandchild:before", + "grandchild:prompt", + "grandchild:after", + "child:after", + "parent:after", + ] + + async def test_resource_template_middleware_three_levels(self): + """Middleware runs at all levels for resource templates.""" + calls: list[str] = [] + + grandchild = FastMCP("Grandchild") + + @grandchild.resource("item://{id}") + 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) as client: + result = await client.read_resource("item://c/gc/42") + assert result[0].text == "item-42" # type: ignore[attr-defined] + + assert calls == [ + "parent:before", + "child:before", + "grandchild:before", + "grandchild:template", + "grandchild:after", + "child:after", + "parent:after", + ] diff --git a/tests/server/proxy/test_stateful_proxy_client.py b/tests/server/proxy/test_stateful_proxy_client.py index e7dc3cb47..e99a5c537 100644 --- a/tests/server/proxy/test_stateful_proxy_client.py +++ b/tests/server/proxy/test_stateful_proxy_client.py @@ -137,8 +137,8 @@ class TestStatefulProxyClient: client_factory=StatefulProxyClient(mcp_b).new_stateful ) multi_proxy_mcp = FastMCP() - multi_proxy_mcp.mount(proxy_mcp_a, prefix="a") - multi_proxy_mcp.mount(proxy_mcp_b, prefix="b") + multi_proxy_mcp.mount(proxy_mcp_a, namespace="a") + multi_proxy_mcp.mount(proxy_mcp_b, namespace="b") async with Client(multi_proxy_mcp) as client: result_a = await client.call_tool("a_tool_a", {}) diff --git a/tests/server/tasks/test_custom_subclass_tasks.py b/tests/server/tasks/test_custom_subclass_tasks.py new file mode 100644 index 000000000..39538836d --- /dev/null +++ b/tests/server/tasks/test_custom_subclass_tasks.py @@ -0,0 +1,184 @@ +"""Tests for custom component subclasses with task support. + +Verifies that custom Tool, Resource, and Prompt subclasses can use +background task execution by setting task_config. +""" + +import asyncio +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from fastmcp import FastMCP +from fastmcp.client import Client +from fastmcp.server.tasks import TaskConfig +from fastmcp.tools.tool import Tool, ToolResult +from fastmcp.utilities.components import FastMCPComponent + + +class CustomTool(Tool): + """A custom tool subclass with task support.""" + + task_config: TaskConfig = TaskConfig(mode="optional") + parameters: dict[str, Any] = {"type": "object", "properties": {}} + + async def run(self, arguments: dict[str, Any]) -> ToolResult: + return ToolResult(content=f"Custom tool executed with {arguments}") + + +class CustomToolWithLogic(Tool): + """A custom tool with actual async work.""" + + task_config: TaskConfig = TaskConfig(mode="optional") + parameters: dict[str, Any] = { + "type": "object", + "properties": {"duration": {"type": "integer"}}, + } + + async def run(self, arguments: dict[str, Any]) -> ToolResult: + duration = arguments.get("duration", 0) + await asyncio.sleep(duration * 0.01) # Short sleep for testing + return ToolResult(content=f"Completed after {duration} units") + + +class CustomToolForbidden(Tool): + """A custom tool with task_config forbidden (default).""" + + parameters: dict[str, Any] = {"type": "object", "properties": {}} + + async def run(self, arguments: dict[str, Any]) -> ToolResult: + return ToolResult(content="Sync only") + + +@pytest.fixture +def custom_tool_server(): + """Create a server with custom tool subclasses.""" + mcp = FastMCP("custom-tool-server") + mcp.add_tool(CustomTool(name="custom_tool", description="A custom tool")) + mcp.add_tool( + CustomToolWithLogic(name="custom_logic", description="Custom tool with logic") + ) + mcp.add_tool( + CustomToolForbidden(name="custom_forbidden", description="No task support") + ) + return mcp + + +async def test_custom_tool_sync_execution(custom_tool_server): + """Custom tool executes synchronously when no task metadata.""" + async with Client(custom_tool_server) as client: + result = await client.call_tool("custom_tool", {}) + assert "Custom tool executed" in str(result) + + +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) as client: + task = await client.call_tool("custom_tool", {}, task=True) + + 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) + + +async def test_custom_tool_with_arguments(custom_tool_server): + """Custom tool receives arguments correctly in background execution.""" + async with Client(custom_tool_server) as client: + task = await client.call_tool("custom_logic", {"duration": 1}, task=True) + + assert task is not None + result = await task.result() + assert "Completed after 1 units" in str(result) + + +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) as client: + # Sync execution works + result = await client.call_tool("custom_forbidden", {}) + assert "Sync only" in str(result) + + +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) as client: + task = await client.call_tool("custom_forbidden", {}, task=True) + + # Should return immediately with error + assert task.returned_immediately + + +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 + + tool = CustomTool(name="test", description="test") + mock_docket = MagicMock() + + tool.register_with_docket(mock_docket) + + # Should register self.run with docket + mock_docket.register.assert_called_once() + call_args = mock_docket.register.call_args + assert call_args[1]["names"] == ["test"] + + +async def test_custom_tool_forbidden_does_not_register(): + """Verify custom tool with forbidden mode doesn't register with docket.""" + tool = CustomToolForbidden(name="test", description="test") + mock_docket = MagicMock() + + tool.register_with_docket(mock_docket) + + # Should NOT register + mock_docket.register.assert_not_called() + + +# ============================================================================== +# Base FastMCPComponent Tests +# ============================================================================== + + +class TestFastMCPComponentDocketMethods: + """Tests for base FastMCPComponent docket integration.""" + + def test_default_task_config_is_forbidden(self): + """Base component defaults to task_config mode='forbidden'.""" + component = FastMCPComponent(name="test") + assert component.task_config.mode == "forbidden" + + def test_register_with_docket_is_noop(self): + """Base register_with_docket does nothing (subclasses override).""" + component = FastMCPComponent(name="test") + mock_docket = MagicMock() + + # Should not raise, just no-op + component.register_with_docket(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'.""" + component = FastMCPComponent(name="test") + mock_docket = MagicMock() + + with pytest.raises(RuntimeError, match="task_config.mode is 'forbidden'"): + await component.add_to_docket(mock_docket) + + async def test_add_to_docket_raises_not_implemented_when_allowed(self): + """Base add_to_docket raises NotImplementedError when not forbidden.""" + component = FastMCPComponent( + name="test", task_config=TaskConfig(mode="optional") + ) + mock_docket = MagicMock() + + with pytest.raises( + NotImplementedError, match="does not implement add_to_docket" + ): + await component.add_to_docket(mock_docket) diff --git a/tests/server/tasks/test_task_mount.py b/tests/server/tasks/test_task_mount.py index bdd234894..70f649732 100644 --- a/tests/server/tasks/test_task_mount.py +++ b/tests/server/tasks/test_task_mount.py @@ -6,14 +6,20 @@ on mounted child servers through a parent server. """ import asyncio +from collections.abc import Sequence +import mcp.types as mt import pytest from docket import Docket from fastmcp import FastMCP from fastmcp.client import Client +from fastmcp.prompts.prompt import PromptResult +from fastmcp.resources.resource import ResourceContent from fastmcp.server.dependencies import CurrentDocket, CurrentFastMCP +from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.server.tasks import TaskConfig +from fastmcp.tools.tool import ToolResult @pytest.fixture(autouse=True) @@ -83,7 +89,7 @@ def parent_server(child_server): return value * 10 # Mount child with prefix - parent.mount(child_server, prefix="child") + parent.mount(child_server, namespace="child") return parent @@ -303,7 +309,7 @@ class TestMountedTaskDependencies: return f"docket available: {docket is not None}" parent = FastMCP("dep-parent") - parent.mount(child, prefix="child") + parent.mount(child, namespace="child") async with Client(parent) as client: task = await client.call_tool("child_tool_with_docket", {}, task=True) @@ -324,7 +330,7 @@ class TestMountedTaskDependencies: return f"server name: {server.name}" parent = FastMCP("server-dep-parent") - parent.mount(child, prefix="child") + parent.mount(child, namespace="child") async with Client(parent) as client: task = await client.call_tool("child_tool_with_server", {}, task=True) @@ -353,8 +359,8 @@ class TestMultipleMounts: return a - b parent = FastMCP("multi-parent") - parent.mount(child1, prefix="math1") - parent.mount(child2, prefix="math2") + parent.mount(child1, namespace="math1") + parent.mount(child2, namespace="math2") async with Client(parent) as client: task1 = await client.call_tool("math1_add", {"a": 10, "b": 5}, task=True) @@ -386,8 +392,8 @@ class TestMountedFunctionNameCollisions: return value * 3 # Triple parent = FastMCP("parent") - parent.mount(child1, prefix="c1") - parent.mount(child2, prefix="c2") + parent.mount(child1, namespace="c1") + parent.mount(child2, namespace="c2") async with Client(parent) as client: # Both should execute their own implementation @@ -433,8 +439,8 @@ class TestMountedFunctionNameCollisions: async def deep_tool() -> str: return "deep" - child.mount(grandchild, prefix="gc") - parent.mount(child, prefix="child") + child.mount(grandchild, namespace="gc") + parent.mount(child, namespace="child") async with Client(parent) as client: # Tool should be accessible and execute correctly @@ -496,7 +502,7 @@ class TestMountedTaskConfigModes: 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, prefix="child") + parent.mount(child_with_modes, namespace="child") return parent async def test_optional_mode_sync_through_mount(self, parent_with_modes): @@ -548,3 +554,223 @@ class TestMountedTaskConfigModes: result = await task.result() # Result is available but may indicate error or sync execution assert result is not None + + +# ----------------------------------------------------------------------------- +# Middleware classes for tracing tests +# ----------------------------------------------------------------------------- + + +class ToolTracingMiddleware(Middleware): + """Middleware that traces tool calls.""" + + def __init__(self, name: str, calls: list[str]): + super().__init__() + self._name = name + self._calls = calls + + async def on_call_tool( + self, + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, ToolResult], + ) -> ToolResult: + self._calls.append(f"{self._name}:before") + result = await call_next(context) + self._calls.append(f"{self._name}:after") + 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, Sequence[ResourceContent]], + ) -> Sequence[ResourceContent]: + 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. + + 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.""" + calls: list[str] = [] + + grandchild = FastMCP("Grandchild") + + @grandchild.tool(task=True) + async def compute(x: int) -> int: + calls.append("grandchild:tool") + 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.mount(child, namespace="c") + parent.add_middleware(ToolTracingMiddleware("parent", calls)) + + async with Client(parent) as client: + task = await client.call_tool("c_gc_compute", {"x": 5}, task=True) + result = await task.result() + assert result.data == 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 + ] + + 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) 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", + ] + + 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) 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", + ] + + 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) 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", + ] diff --git a/tests/server/test_mount.py b/tests/server/test_mount.py index 947171b6b..615288253 100644 --- a/tests/server/test_mount.py +++ b/tests/server/test_mount.py @@ -77,7 +77,7 @@ class TestBasicMount: return "This is from the sub app" # Mount with empty prefix but without deprecated separators - main_app.mount(sub_app, prefix=prefix) + main_app.mount(sub_app, namespace=prefix) tools = await main_app.get_tools() # With empty prefix, the tool should keep its original name @@ -1222,8 +1222,8 @@ class TestDeeplyNestedMount: def multiply(a: int, b: int) -> int: return a * b - middle.mount(leaf, prefix="leaf") - root.mount(middle, prefix="middle") + middle.mount(leaf, namespace="leaf") + root.mount(middle, namespace="middle") async with Client(root) as client: # Tool at level 2 should work @@ -1248,8 +1248,8 @@ class TestDeeplyNestedMount: def middle_data() -> str: return "middle data" - middle.mount(leaf, prefix="leaf") - root.mount(middle, prefix="middle") + middle.mount(leaf, namespace="leaf") + root.mount(middle, namespace="middle") async with Client(root) as client: # Resource at level 2 should work @@ -1274,8 +1274,8 @@ class TestDeeplyNestedMount: def middle_item(id: str) -> str: return f"middle item {id}" - middle.mount(leaf, prefix="leaf") - root.mount(middle, prefix="middle") + middle.mount(leaf, namespace="leaf") + root.mount(middle, namespace="middle") async with Client(root) as client: # Resource template at level 2 should work @@ -1300,8 +1300,8 @@ class TestDeeplyNestedMount: def middle_prompt(name: str) -> str: return f"Hello from middle: {name}" - middle.mount(leaf, prefix="leaf") - root.mount(middle, prefix="middle") + middle.mount(leaf, namespace="leaf") + root.mount(middle, namespace="middle") async with Client(root) as client: # Prompt at level 2 should work @@ -1325,9 +1325,9 @@ class TestDeeplyNestedMount: def deep_tool() -> str: return "very deep" - level2.mount(level3, prefix="l3") - level1.mount(level2, prefix="l2") - root.mount(level1, prefix="l1") + level2.mount(level3, namespace="l3") + level1.mount(level2, namespace="l2") + root.mount(level1, namespace="l1") async with Client(root) as client: # Verify tool is listed diff --git a/tests/server/test_providers.py b/tests/server/test_providers.py index dee1f3273..5c4761053 100644 --- a/tests/server/test_providers.py +++ b/tests/server/test_providers.py @@ -4,14 +4,14 @@ from collections.abc import Sequence from typing import Any import pytest -from mcp.types import AnyUrl, PromptMessage, TextContent +from mcp.types import AnyUrl from mcp.types import Tool as MCPTool from fastmcp import FastMCP from fastmcp.client import Client from fastmcp.client.client import CallToolResult -from fastmcp.prompts.prompt import FunctionPrompt, Prompt, PromptResult -from fastmcp.resources.resource import FunctionResource, Resource, ResourceContent +from fastmcp.prompts.prompt import FunctionPrompt, Prompt +from fastmcp.resources.resource import FunctionResource, Resource from fastmcp.resources.template import FunctionResourceTemplate, ResourceTemplate from fastmcp.server.providers import Provider from fastmcp.tools.tool import Tool, ToolResult @@ -213,12 +213,11 @@ class TestProvider: async with Client(base_server) as client: await client.call_tool(name="dynamic_multiply", arguments={"a": 2, "b": 3}) - # get_tool is called three times: - # 1. Server.get_tool() for task config check calls provider.get_tool() - # 2. _call_tool() calls provider.get_tool() to check _should_enable_component - # 3. Default call_tool() implementation calls get_tool() internally + # get_tool is called once for efficient lookup: + # _call_tool() calls provider.get_tool() to get the tool and execute it + # (task config is checked inside the tool's _run() method, not via a separate lookup) # Key point: list_tools is NOT called during tool execution (efficient lookup) - assert provider.get_tool_call_count == 3 + assert provider.get_tool_call_count == 1 async def test_default_get_tool_falls_back_to_list(self, base_server: FastMCP): """Test that BaseToolProvider's default get_tool calls list_tools.""" @@ -382,55 +381,6 @@ class TestProviderExecutionMethods: assert result.structured_content is not None assert result.structured_content["result"] == 3 # type: ignore[attr-defined] - async def test_call_tool_custom_implementation(self): - """Test that providers can override call_tool for custom behavior.""" - - class CustomCallProvider(Provider): - """Provider that wraps tool execution with custom logic.""" - - def __init__(self): - super().__init__() - self.call_count = 0 - self._tool = SimpleTool( - name="custom_tool", - description="Test", - parameters={"type": "object", "properties": {"a": {}, "b": {}}}, - operation="add", - ) - - async def list_tools(self) -> Sequence[Tool]: - return [self._tool] - - async def get_tool(self, name: str) -> Tool | None: - if name == "custom_tool": - return self._tool - return None - - async def call_tool( - self, name: str, arguments: dict[str, Any] - ) -> ToolResult | None: - # Custom behavior: track calls and modify result - self.call_count += 1 - tool = await self.get_tool(name) - if tool is None: - return None - result = await tool.run(arguments) - # Add custom metadata to result - result.structured_content["custom_wrapper"] = True # type: ignore[index] - return result - - provider = CustomCallProvider() - mcp = FastMCP("TestServer") - mcp.add_provider(provider) - - async with Client(mcp) as client: - result = await client.call_tool("custom_tool", {"a": 5, "b": 3}) - - assert provider.call_count == 1 - assert result.structured_content is not None - assert result.structured_content["result"] == 8 # type: ignore[attr-defined] - assert result.structured_content["custom_wrapper"] is True # type: ignore[attr-defined] - async def test_read_resource_default_implementation(self): """Test that default read_resource uses get_resource and reads it.""" @@ -454,37 +404,6 @@ class TestProviderExecutionMethods: assert len(result) == 1 assert result[0].text == "hello world" - async def test_read_resource_custom_implementation(self): - """Test that providers can override read_resource for custom behavior.""" - - class CustomReadProvider(Provider): - """Provider that transforms resource content.""" - - async def list_resources(self) -> Sequence[Resource]: - return [ - FunctionResource( - uri=AnyUrl("test://data"), - name="Test Data", - fn=lambda: "original", - ) - ] - - async def read_resource(self, uri: str) -> ResourceContent | None: - if uri == "test://data": - # Custom behavior: return transformed content - return ResourceContent(content="TRANSFORMED") - return None - - provider = CustomReadProvider() - mcp = FastMCP("TestServer") - mcp.add_provider(provider) - - async with Client(mcp) as client: - result = await client.read_resource("test://data") - - assert len(result) == 1 - assert result[0].text == "TRANSFORMED" - async def test_read_resource_template_default(self): """Test that read_resource_template handles template-based resources.""" @@ -530,46 +449,3 @@ class TestProviderExecutionMethods: assert len(result.messages) == 1 assert result.messages[0].content.text == "Hello, World!" # type: ignore[attr-defined] - - async def test_render_prompt_custom_implementation(self): - """Test that providers can override render_prompt for custom behavior.""" - - class CustomRenderProvider(Provider): - """Provider that adds prefix to all prompts.""" - - async def list_prompts(self) -> Sequence[Prompt]: - return [ - FunctionPrompt.from_function( - fn=lambda: "original message", - name="test_prompt", - description="Test", - ) - ] - - async def render_prompt( - self, name: str, arguments: dict[str, Any] | None - ) -> PromptResult | None: - if name == "test_prompt": - # Custom behavior: add prefix - return PromptResult( - messages=[ - PromptMessage( - role="user", - content=TextContent( - type="text", - text="[CUSTOM PREFIX] original message", - ), - ) - ] - ) - return None - - provider = CustomRenderProvider() - mcp = FastMCP("TestServer") - mcp.add_provider(provider) - - async with Client(mcp) as client: - result = await client.get_prompt("test_prompt", {}) - - assert len(result.messages) == 1 - assert "[CUSTOM PREFIX]" in result.messages[0].content.text # type: ignore[attr-defined] diff --git a/tests/tools/test_tool_manager.py b/tests/tools/test_tool_manager.py index 9843f5d75..c9c88fd98 100644 --- a/tests/tools/test_tool_manager.py +++ b/tests/tools/test_tool_manager.py @@ -1035,7 +1035,7 @@ class TestMountedComponentsRaiseOnLoadError: child_mcp = FastMCP("FailingChildServer") # Create a failing mounted server by corrupting it - parent_mcp.mount(child_mcp, prefix="child") + parent_mcp.mount(child_mcp, namespace="child") # Corrupt the parent's providers to make it fail during loading parent_mcp._providers.append("invalid") # type: ignore @@ -1049,7 +1049,7 @@ class TestMountedComponentsRaiseOnLoadError: child_mcp = FastMCP("FailingChildServer") # Create a failing mounted server - parent_mcp.mount(child_mcp, prefix="child") + parent_mcp.mount(child_mcp, namespace="child") # Corrupt the parent's providers to make it fail during loading parent_mcp._providers.append("invalid") # type: ignore