Align prompt handler with resource pattern (#2740)

This commit is contained in:
Jeremiah Lowin 2025-12-25 22:47:50 -05:00 committed by GitHub
commit 5a6f71503b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 41 additions and 39 deletions

View file

@ -24,9 +24,6 @@ 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.server.dependencies import _docket_fn_key, _task_metadata
from fastmcp.utilities.logging import get_logger
if TYPE_CHECKING:
@ -254,20 +251,19 @@ class LowLevelServer(_Server[LifespanResultT, RequestT]):
[
Callable[
[str, dict[str, Any] | None],
Awaitable[PromptResult | mcp.types.CreateTaskResult],
Awaitable[mcp.types.GetPromptResult | mcp.types.CreateTaskResult],
]
],
Callable[
[str, dict[str, Any] | None],
Awaitable[PromptResult | mcp.types.CreateTaskResult],
Awaitable[mcp.types.GetPromptResult | 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.
for background task execution. This decorator wraps the result in ServerResult.
This decorator can be removed once the MCP SDK adds native CreateTaskResult support
for prompts.
@ -276,41 +272,17 @@ class LowLevelServer(_Server[LifespanResultT, RequestT]):
def decorator(
func: Callable[
[str, dict[str, Any] | None],
Awaitable[PromptResult | mcp.types.CreateTaskResult],
Awaitable[mcp.types.GetPromptResult | mcp.types.CreateTaskResult],
],
) -> Callable[
[str, dict[str, Any] | None],
Awaitable[PromptResult | mcp.types.CreateTaskResult],
Awaitable[mcp.types.GetPromptResult | 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)
result = await func(req.params.name, req.params.arguments)
return mcp.types.ServerResult(result)
self.request_handlers[mcp.types.GetPromptRequest] = handler
return func

View file

@ -1584,18 +1584,48 @@ class FastMCP(Generic[LifespanResultT]):
async def _get_prompt_mcp(
self, name: str, arguments: dict[str, Any] | None
) -> PromptResult | mcp.types.CreateTaskResult:
) -> mcp.types.GetPromptResult | mcp.types.CreateTaskResult:
"""Handle MCP 'getPrompt' requests.
The LowLevelServer.get_prompt() decorator handles task metadata,
contextvars, and MCP conversion.
Sets task metadata contextvar and calls render_prompt(). The prompt's
_render() method handles the backgrounding decision.
Args:
name: The prompt name
arguments: Prompt arguments
Returns:
GetPromptResult or CreateTaskResult for background execution
"""
from fastmcp.server.dependencies import _docket_fn_key, _task_metadata
logger.debug(
f"[{self.name}] Handler called: get_prompt %s with %s", name, arguments
)
try:
return await self.render_prompt(name, arguments)
# Extract SEP-1686 task metadata from 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
# 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 result
return 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: