mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
Consistent decorator-based MCP handler registration (#2732)
This commit is contained in:
parent
9f72096c67
commit
9f5a177126
3 changed files with 183 additions and 142 deletions
|
|
@ -303,7 +303,12 @@ class Context:
|
|||
Returns:
|
||||
The prompt result
|
||||
"""
|
||||
return await self.fastmcp._get_prompt_mcp(name, arguments)
|
||||
result = await self.fastmcp.render_prompt(name, arguments)
|
||||
if isinstance(result, mcp.types.CreateTaskResult):
|
||||
raise RuntimeError(
|
||||
"Unexpected CreateTaskResult: Context calls should not have task metadata"
|
||||
)
|
||||
return result.to_mcp_prompt_result()
|
||||
|
||||
async def read_resource(self, uri: str | AnyUrl) -> list[ResourceContent]:
|
||||
"""Read a resource by URI.
|
||||
|
|
@ -314,8 +319,12 @@ class Context:
|
|||
Returns:
|
||||
List of ResourceContent objects
|
||||
"""
|
||||
# Context calls don't have task metadata, so always returns list
|
||||
return await self.fastmcp._read_resource_mcp(uri)
|
||||
result = await self.fastmcp.read_resource(str(uri))
|
||||
if isinstance(result, mcp.types.CreateTaskResult):
|
||||
raise RuntimeError(
|
||||
"Unexpected CreateTaskResult: Context calls should not have task metadata"
|
||||
)
|
||||
return result
|
||||
|
||||
async def log(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import AsyncExitStack
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
|
|
@ -21,7 +22,13 @@ from mcp.server.session import ServerSession
|
|||
from mcp.server.stdio import stdio_server as stdio_server
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import RequestResponder
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from fastmcp.prompts import Prompt
|
||||
from fastmcp.prompts.prompt import PromptResult
|
||||
from fastmcp.resources import Resource
|
||||
from fastmcp.resources.types import ResourceContent
|
||||
from fastmcp.server.dependencies import _docket_fn_key, _task_metadata
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -198,3 +205,143 @@ class LowLevelServer(_Server[LifespanResultT, RequestT]):
|
|||
lifespan_context,
|
||||
raise_exceptions,
|
||||
)
|
||||
|
||||
def read_resource(
|
||||
self,
|
||||
) -> Callable[
|
||||
[
|
||||
Callable[
|
||||
[AnyUrl],
|
||||
Awaitable[list[ResourceContent] | mcp.types.CreateTaskResult],
|
||||
]
|
||||
],
|
||||
Callable[
|
||||
[AnyUrl],
|
||||
Awaitable[list[ResourceContent] | mcp.types.CreateTaskResult],
|
||||
],
|
||||
]:
|
||||
"""
|
||||
Decorator for registering a read_resource handler with CreateTaskResult support.
|
||||
|
||||
The MCP SDK's read_resource decorator does not support returning CreateTaskResult
|
||||
for background task execution. This decorator provides that support by wrapping the
|
||||
handler with task metadata extraction, contextvar management, and MCP format conversion.
|
||||
|
||||
This decorator can be removed once the MCP SDK adds native CreateTaskResult support
|
||||
for resources.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[
|
||||
[AnyUrl],
|
||||
Awaitable[list[ResourceContent] | mcp.types.CreateTaskResult],
|
||||
],
|
||||
) -> Callable[
|
||||
[AnyUrl],
|
||||
Awaitable[list[ResourceContent] | mcp.types.CreateTaskResult],
|
||||
]:
|
||||
async def handler(
|
||||
req: mcp.types.ReadResourceRequest,
|
||||
) -> mcp.types.ServerResult:
|
||||
uri = req.params.uri
|
||||
|
||||
# Extract task metadata from request context
|
||||
task_meta_dict: dict[str, Any] | None = None
|
||||
try:
|
||||
ctx = self.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
|
||||
|
||||
# Set contextvars
|
||||
task_token = _task_metadata.set(task_meta_dict)
|
||||
key_token = _docket_fn_key.set(Resource.make_key(str(uri)))
|
||||
try:
|
||||
result = await func(uri)
|
||||
|
||||
if isinstance(result, mcp.types.CreateTaskResult):
|
||||
return mcp.types.ServerResult(result)
|
||||
|
||||
contents = [item.to_mcp_resource_contents(uri) for item in result]
|
||||
return mcp.types.ServerResult(
|
||||
mcp.types.ReadResourceResult(contents=contents)
|
||||
)
|
||||
finally:
|
||||
_task_metadata.reset(task_token)
|
||||
_docket_fn_key.reset(key_token)
|
||||
|
||||
self.request_handlers[mcp.types.ReadResourceRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def get_prompt(
|
||||
self,
|
||||
) -> Callable[
|
||||
[
|
||||
Callable[
|
||||
[str, dict[str, Any] | None],
|
||||
Awaitable[PromptResult | mcp.types.CreateTaskResult],
|
||||
]
|
||||
],
|
||||
Callable[
|
||||
[str, dict[str, Any] | None],
|
||||
Awaitable[PromptResult | mcp.types.CreateTaskResult],
|
||||
],
|
||||
]:
|
||||
"""
|
||||
Decorator for registering a get_prompt handler with CreateTaskResult support.
|
||||
|
||||
The MCP SDK's get_prompt decorator does not support returning CreateTaskResult
|
||||
for background task execution. This decorator provides that support by wrapping the
|
||||
handler with task metadata extraction, contextvar management, and MCP format conversion.
|
||||
|
||||
This decorator can be removed once the MCP SDK adds native CreateTaskResult support
|
||||
for prompts.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[
|
||||
[str, dict[str, Any] | None],
|
||||
Awaitable[PromptResult | mcp.types.CreateTaskResult],
|
||||
],
|
||||
) -> Callable[
|
||||
[str, dict[str, Any] | None],
|
||||
Awaitable[PromptResult | mcp.types.CreateTaskResult],
|
||||
]:
|
||||
async def handler(
|
||||
req: mcp.types.GetPromptRequest,
|
||||
) -> mcp.types.ServerResult:
|
||||
name = req.params.name
|
||||
arguments = req.params.arguments
|
||||
|
||||
# Extract task metadata from request context
|
||||
task_meta_dict: dict[str, Any] | None = None
|
||||
try:
|
||||
ctx = self.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
|
||||
|
||||
# Set contextvars
|
||||
task_token = _task_metadata.set(task_meta_dict)
|
||||
key_token = _docket_fn_key.set(Prompt.make_key(name))
|
||||
try:
|
||||
result = await func(name, arguments)
|
||||
|
||||
if isinstance(result, mcp.types.CreateTaskResult):
|
||||
return mcp.types.ServerResult(result)
|
||||
|
||||
return mcp.types.ServerResult(result.to_mcp_prompt_result())
|
||||
finally:
|
||||
_task_metadata.reset(task_token)
|
||||
_docket_fn_key.reset(key_token)
|
||||
|
||||
self.request_handlers[mcp.types.GetPromptRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
|
|
|||
|
|
@ -38,7 +38,6 @@ from mcp.types import (
|
|||
AnyFunction,
|
||||
CallToolRequestParams,
|
||||
ContentBlock,
|
||||
GetPromptResult,
|
||||
ToolAnnotations,
|
||||
)
|
||||
from mcp.types import Prompt as SDKPrompt
|
||||
|
|
@ -608,34 +607,21 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
def _setup_handlers(self) -> None:
|
||||
"""Set up core MCP protocol handlers.
|
||||
|
||||
We override the SDK's default handlers for tools/call, resources/read,
|
||||
and prompts/get to add task-augmented execution support (SEP-1686).
|
||||
|
||||
The SDK's decorators have different capabilities:
|
||||
- call_tool: Supports CreateTaskResult returns AND validate_input
|
||||
- read_resource: Does NOT support CreateTaskResult
|
||||
- get_prompt: Does NOT support CreateTaskResult
|
||||
|
||||
So we use the SDK decorator for tools (to get input validation), but
|
||||
register custom handlers for resources and prompts.
|
||||
All handlers use decorator-based registration for consistency.
|
||||
The call_tool decorator is from the SDK (supports CreateTaskResult + validate_input).
|
||||
The read_resource and get_prompt decorators are from LowLevelServer to add
|
||||
CreateTaskResult support until the SDK provides it natively.
|
||||
"""
|
||||
self._mcp_server.list_tools()(self._list_tools_mcp)
|
||||
self._mcp_server.list_resources()(self._list_resources_mcp)
|
||||
self._mcp_server.list_resource_templates()(self._list_resource_templates_mcp)
|
||||
self._mcp_server.list_prompts()(self._list_prompts_mcp)
|
||||
|
||||
# Tools: SDK decorator provides validate_input + CreateTaskResult support
|
||||
self._mcp_server.call_tool(validate_input=self.strict_input_validation)(
|
||||
self._call_tool_mcp
|
||||
)
|
||||
|
||||
# Resources/Prompts: Custom handlers (SDK decorators don't support CreateTaskResult)
|
||||
self._mcp_server.request_handlers[mcp.types.ReadResourceRequest] = (
|
||||
self._read_resource_handler
|
||||
)
|
||||
self._mcp_server.request_handlers[mcp.types.GetPromptRequest] = (
|
||||
self._get_prompt_handler
|
||||
)
|
||||
self._mcp_server.read_resource()(self._read_resource_mcp)
|
||||
self._mcp_server.get_prompt()(self._get_prompt_mcp)
|
||||
|
||||
# Register SEP-1686 task protocol handlers
|
||||
self._setup_task_protocol_handlers()
|
||||
|
|
@ -1552,142 +1538,41 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
except NotFoundError as e:
|
||||
raise NotFoundError(f"Unknown tool: {key!r}") from e
|
||||
|
||||
async def _read_resource_handler(
|
||||
self, req: mcp.types.ReadResourceRequest
|
||||
) -> mcp.types.ServerResult:
|
||||
"""Handle resources/read requests with task-augmented execution support.
|
||||
async def _read_resource_mcp(
|
||||
self, uri: AnyUrl | str
|
||||
) -> list[ResourceContent] | mcp.types.CreateTaskResult:
|
||||
"""Handle MCP 'readResource' requests.
|
||||
|
||||
This is a custom handler because the SDK's read_resource decorator
|
||||
does not support returning CreateTaskResult for background tasks.
|
||||
"""
|
||||
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_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
|
||||
|
||||
try:
|
||||
# Set contextvars so Resource._read() can access them
|
||||
task_token = _task_metadata.set(task_meta_dict)
|
||||
key_token = _docket_fn_key.set(Resource.make_key(str(uri)))
|
||||
try:
|
||||
result = await self.read_resource(str(uri))
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
) -> mcp.types.ServerResult:
|
||||
"""Handle prompts/get requests with task-augmented execution support.
|
||||
|
||||
This is a custom handler because the SDK's get_prompt decorator
|
||||
does not support returning CreateTaskResult for background tasks.
|
||||
"""
|
||||
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_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
|
||||
|
||||
try:
|
||||
# Set contextvars so Prompt._render() can access them
|
||||
task_token = _task_metadata.set(task_meta_dict)
|
||||
key_token = _docket_fn_key.set(Prompt.make_key(name))
|
||||
try:
|
||||
result = await self.render_prompt(name, arguments)
|
||||
|
||||
if isinstance(result, mcp.types.CreateTaskResult):
|
||||
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)
|
||||
|
||||
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 _read_resource_mcp(self, uri: AnyUrl | str) -> list[ResourceContent]:
|
||||
"""
|
||||
Handle MCP 'readResource' requests (used by Context.read_resource()).
|
||||
|
||||
Delegates to read_resource() without task metadata, so CreateTaskResult
|
||||
should never be returned.
|
||||
The LowLevelServer.read_resource() decorator handles task metadata,
|
||||
contextvars, and MCP conversion.
|
||||
"""
|
||||
logger.debug(f"[{self.name}] Handler called: read_resource %s", uri)
|
||||
|
||||
try:
|
||||
result = await self.read_resource(str(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
|
||||
return await self.read_resource(str(uri))
|
||||
except DisabledError as e:
|
||||
# convert to NotFoundError to avoid leaking resource presence
|
||||
raise NotFoundError(f"Unknown resource: {str(uri)!r}") from e
|
||||
except NotFoundError as e:
|
||||
# standardize NotFound message
|
||||
raise NotFoundError(f"Unknown resource: {str(uri)!r}") from e
|
||||
except NotFoundError:
|
||||
raise
|
||||
|
||||
async def _get_prompt_mcp(
|
||||
self, name: str, arguments: dict[str, Any] | None = None
|
||||
) -> GetPromptResult:
|
||||
"""
|
||||
Handle MCP 'getPrompt' requests (used by Context.get_prompt()).
|
||||
self, name: str, arguments: dict[str, Any] | None
|
||||
) -> PromptResult | mcp.types.CreateTaskResult:
|
||||
"""Handle MCP 'getPrompt' requests.
|
||||
|
||||
Delegates to render_prompt() and converts to MCP SDK type.
|
||||
The LowLevelServer.get_prompt() decorator handles task metadata,
|
||||
contextvars, and MCP conversion.
|
||||
"""
|
||||
logger.debug(
|
||||
f"[{self.name}] Handler called: get_prompt %s with %s", name, arguments
|
||||
)
|
||||
|
||||
try:
|
||||
result = await self.render_prompt(name, arguments)
|
||||
if isinstance(result, mcp.types.CreateTaskResult):
|
||||
# Should never happen without task metadata
|
||||
raise RuntimeError("Unexpected CreateTaskResult in _get_prompt_mcp")
|
||||
return result.to_mcp_prompt_result()
|
||||
return await self.render_prompt(name, arguments)
|
||||
except DisabledError as e:
|
||||
# convert to NotFoundError to avoid leaking prompt presence
|
||||
raise NotFoundError(f"Unknown prompt: {name}") from e
|
||||
except NotFoundError as e:
|
||||
# standardize NotFound message
|
||||
raise NotFoundError(f"Unknown prompt: {name}") from e
|
||||
raise NotFoundError(f"Unknown prompt: {name!r}") from e
|
||||
except NotFoundError:
|
||||
raise
|
||||
|
||||
def add_tool(self, tool: Tool) -> Tool:
|
||||
"""Add a tool to the server.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue