SEP-1686: Tasks (#1645)
This commit is contained in:
parent
5983a650cc
commit
c92bb2f7ff
2 changed files with 14230 additions and 120 deletions
288
src/mcp/server/lowlevel/experimental.py
Normal file
288
src/mcp/server/lowlevel/experimental.py
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
"""Experimental handlers for the low-level MCP server.
|
||||
|
||||
WARNING: These APIs are experimental and may change without notice.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from mcp.server.experimental.task_support import TaskSupport
|
||||
from mcp.server.lowlevel.func_inspection import create_call_wrapper
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.experimental.tasks.helpers import cancel_task
|
||||
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
|
||||
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
|
||||
from mcp.shared.experimental.tasks.store import TaskStore
|
||||
from mcp.types import (
|
||||
INVALID_PARAMS,
|
||||
CancelTaskRequest,
|
||||
CancelTaskResult,
|
||||
ErrorData,
|
||||
GetTaskPayloadRequest,
|
||||
GetTaskPayloadResult,
|
||||
GetTaskRequest,
|
||||
GetTaskResult,
|
||||
ListTasksRequest,
|
||||
ListTasksResult,
|
||||
ServerCapabilities,
|
||||
ServerResult,
|
||||
ServerTasksCapability,
|
||||
ServerTasksRequestsCapability,
|
||||
TasksCancelCapability,
|
||||
TasksListCapability,
|
||||
TasksToolsCapability,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.server.lowlevel.server import Server
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ExperimentalHandlers:
|
||||
"""Experimental request/notification handlers.
|
||||
|
||||
WARNING: These APIs are experimental and may change without notice.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server: Server,
|
||||
request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]],
|
||||
notification_handlers: dict[type, Callable[..., Awaitable[None]]],
|
||||
):
|
||||
self._server = server
|
||||
self._request_handlers = request_handlers
|
||||
self._notification_handlers = notification_handlers
|
||||
self._task_support: TaskSupport | None = None
|
||||
|
||||
@property
|
||||
def task_support(self) -> TaskSupport | None:
|
||||
"""Get the task support configuration, if enabled."""
|
||||
return self._task_support
|
||||
|
||||
def update_capabilities(self, capabilities: ServerCapabilities) -> None:
|
||||
# Only add tasks capability if handlers are registered
|
||||
if not any(
|
||||
req_type in self._request_handlers
|
||||
for req_type in [GetTaskRequest, ListTasksRequest, CancelTaskRequest, GetTaskPayloadRequest]
|
||||
):
|
||||
return
|
||||
|
||||
capabilities.tasks = ServerTasksCapability()
|
||||
if ListTasksRequest in self._request_handlers:
|
||||
capabilities.tasks.list = TasksListCapability()
|
||||
if CancelTaskRequest in self._request_handlers:
|
||||
capabilities.tasks.cancel = TasksCancelCapability()
|
||||
|
||||
capabilities.tasks.requests = ServerTasksRequestsCapability(
|
||||
tools=TasksToolsCapability()
|
||||
) # assuming always supported for now
|
||||
|
||||
def enable_tasks(
|
||||
self,
|
||||
store: TaskStore | None = None,
|
||||
queue: TaskMessageQueue | None = None,
|
||||
) -> TaskSupport:
|
||||
"""
|
||||
Enable experimental task support.
|
||||
|
||||
This sets up the task infrastructure and auto-registers default handlers
|
||||
for tasks/get, tasks/result, tasks/list, and tasks/cancel.
|
||||
|
||||
Args:
|
||||
store: Custom TaskStore implementation (defaults to InMemoryTaskStore)
|
||||
queue: Custom TaskMessageQueue implementation (defaults to InMemoryTaskMessageQueue)
|
||||
|
||||
Returns:
|
||||
The TaskSupport configuration object
|
||||
|
||||
Example:
|
||||
# Simple in-memory setup
|
||||
server.experimental.enable_tasks()
|
||||
|
||||
# Custom store/queue for distributed systems
|
||||
server.experimental.enable_tasks(
|
||||
store=RedisTaskStore(redis_url),
|
||||
queue=RedisTaskMessageQueue(redis_url),
|
||||
)
|
||||
|
||||
WARNING: This API is experimental and may change without notice.
|
||||
"""
|
||||
if store is None:
|
||||
store = InMemoryTaskStore()
|
||||
if queue is None:
|
||||
queue = InMemoryTaskMessageQueue()
|
||||
|
||||
self._task_support = TaskSupport(store=store, queue=queue)
|
||||
|
||||
# Auto-register default handlers
|
||||
self._register_default_task_handlers()
|
||||
|
||||
return self._task_support
|
||||
|
||||
def _register_default_task_handlers(self) -> None:
|
||||
"""Register default handlers for task operations."""
|
||||
assert self._task_support is not None
|
||||
support = self._task_support
|
||||
|
||||
# Register get_task handler if not already registered
|
||||
if GetTaskRequest not in self._request_handlers:
|
||||
|
||||
async def _default_get_task(req: GetTaskRequest) -> ServerResult:
|
||||
task = await support.store.get_task(req.params.taskId)
|
||||
if task is None:
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
code=INVALID_PARAMS,
|
||||
message=f"Task not found: {req.params.taskId}",
|
||||
)
|
||||
)
|
||||
return ServerResult(
|
||||
GetTaskResult(
|
||||
taskId=task.taskId,
|
||||
status=task.status,
|
||||
statusMessage=task.statusMessage,
|
||||
createdAt=task.createdAt,
|
||||
lastUpdatedAt=task.lastUpdatedAt,
|
||||
ttl=task.ttl,
|
||||
pollInterval=task.pollInterval,
|
||||
)
|
||||
)
|
||||
|
||||
self._request_handlers[GetTaskRequest] = _default_get_task
|
||||
|
||||
# Register get_task_result handler if not already registered
|
||||
if GetTaskPayloadRequest not in self._request_handlers:
|
||||
|
||||
async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult:
|
||||
ctx = self._server.request_context
|
||||
result = await support.handler.handle(req, ctx.session, ctx.request_id)
|
||||
return ServerResult(result)
|
||||
|
||||
self._request_handlers[GetTaskPayloadRequest] = _default_get_task_result
|
||||
|
||||
# Register list_tasks handler if not already registered
|
||||
if ListTasksRequest not in self._request_handlers:
|
||||
|
||||
async def _default_list_tasks(req: ListTasksRequest) -> ServerResult:
|
||||
cursor = req.params.cursor if req.params else None
|
||||
tasks, next_cursor = await support.store.list_tasks(cursor)
|
||||
return ServerResult(ListTasksResult(tasks=tasks, nextCursor=next_cursor))
|
||||
|
||||
self._request_handlers[ListTasksRequest] = _default_list_tasks
|
||||
|
||||
# Register cancel_task handler if not already registered
|
||||
if CancelTaskRequest not in self._request_handlers:
|
||||
|
||||
async def _default_cancel_task(req: CancelTaskRequest) -> ServerResult:
|
||||
result = await cancel_task(support.store, req.params.taskId)
|
||||
return ServerResult(result)
|
||||
|
||||
self._request_handlers[CancelTaskRequest] = _default_cancel_task
|
||||
|
||||
def list_tasks(
|
||||
self,
|
||||
) -> Callable[
|
||||
[Callable[[ListTasksRequest], Awaitable[ListTasksResult]]],
|
||||
Callable[[ListTasksRequest], Awaitable[ListTasksResult]],
|
||||
]:
|
||||
"""Register a handler for listing tasks.
|
||||
|
||||
WARNING: This API is experimental and may change without notice.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[[ListTasksRequest], Awaitable[ListTasksResult]],
|
||||
) -> Callable[[ListTasksRequest], Awaitable[ListTasksResult]]:
|
||||
logger.debug("Registering handler for ListTasksRequest")
|
||||
wrapper = create_call_wrapper(func, ListTasksRequest)
|
||||
|
||||
async def handler(req: ListTasksRequest) -> ServerResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
|
||||
self._request_handlers[ListTasksRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def get_task(
|
||||
self,
|
||||
) -> Callable[
|
||||
[Callable[[GetTaskRequest], Awaitable[GetTaskResult]]], Callable[[GetTaskRequest], Awaitable[GetTaskResult]]
|
||||
]:
|
||||
"""Register a handler for getting task status.
|
||||
|
||||
WARNING: This API is experimental and may change without notice.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]],
|
||||
) -> Callable[[GetTaskRequest], Awaitable[GetTaskResult]]:
|
||||
logger.debug("Registering handler for GetTaskRequest")
|
||||
wrapper = create_call_wrapper(func, GetTaskRequest)
|
||||
|
||||
async def handler(req: GetTaskRequest) -> ServerResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
|
||||
self._request_handlers[GetTaskRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def get_task_result(
|
||||
self,
|
||||
) -> Callable[
|
||||
[Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]],
|
||||
Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
|
||||
]:
|
||||
"""Register a handler for getting task results/payload.
|
||||
|
||||
WARNING: This API is experimental and may change without notice.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]],
|
||||
) -> Callable[[GetTaskPayloadRequest], Awaitable[GetTaskPayloadResult]]:
|
||||
logger.debug("Registering handler for GetTaskPayloadRequest")
|
||||
wrapper = create_call_wrapper(func, GetTaskPayloadRequest)
|
||||
|
||||
async def handler(req: GetTaskPayloadRequest) -> ServerResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
|
||||
self._request_handlers[GetTaskPayloadRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def cancel_task(
|
||||
self,
|
||||
) -> Callable[
|
||||
[Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]],
|
||||
Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
|
||||
]:
|
||||
"""Register a handler for cancelling tasks.
|
||||
|
||||
WARNING: This API is experimental and may change without notice.
|
||||
"""
|
||||
|
||||
def decorator(
|
||||
func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]],
|
||||
) -> Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]:
|
||||
logger.debug("Registering handler for CancelTaskRequest")
|
||||
wrapper = create_call_wrapper(func, CancelTaskRequest)
|
||||
|
||||
async def handler(req: CancelTaskRequest) -> ServerResult:
|
||||
result = await wrapper(req)
|
||||
return ServerResult(result)
|
||||
|
||||
self._request_handlers[CancelTaskRequest] = handler
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
|
@ -67,6 +67,7 @@ messages from the client.
|
|||
|
||||
from __future__ import annotations as _annotations
|
||||
|
||||
import base64
|
||||
import contextvars
|
||||
import json
|
||||
import logging
|
||||
|
|
@ -82,6 +83,8 @@ from pydantic import AnyUrl
|
|||
from typing_extensions import TypeVar
|
||||
|
||||
import mcp.types as types
|
||||
from mcp.server.experimental.request_context import Experimental
|
||||
from mcp.server.lowlevel.experimental import ExperimentalHandlers
|
||||
from mcp.server.lowlevel.func_inspection import create_call_wrapper
|
||||
from mcp.server.lowlevel.helper_types import ReadResourceContents
|
||||
from mcp.server.models import InitializationOptions
|
||||
|
|
@ -155,6 +158,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
}
|
||||
self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
|
||||
self._tool_cache: dict[str, types.Tool] = {}
|
||||
self._experimental_handlers: ExperimentalHandlers | None = None
|
||||
logger.debug("Initializing server %r", name)
|
||||
|
||||
def create_initialization_options(
|
||||
|
|
@ -220,7 +224,7 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
if types.CompleteRequest in self.request_handlers:
|
||||
completions_capability = types.CompletionsCapability()
|
||||
|
||||
return types.ServerCapabilities(
|
||||
capabilities = types.ServerCapabilities(
|
||||
prompts=prompts_capability,
|
||||
resources=resources_capability,
|
||||
tools=tools_capability,
|
||||
|
|
@ -228,6 +232,9 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
experimental=experimental_capabilities,
|
||||
completions=completions_capability,
|
||||
)
|
||||
if self._experimental_handlers:
|
||||
self._experimental_handlers.update_capabilities(capabilities)
|
||||
return capabilities
|
||||
|
||||
@property
|
||||
def request_context(
|
||||
|
|
@ -236,6 +243,18 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
"""If called outside of a request context, this will raise a LookupError."""
|
||||
return request_ctx.get()
|
||||
|
||||
@property
|
||||
def experimental(self) -> ExperimentalHandlers:
|
||||
"""Experimental APIs for tasks and other features.
|
||||
|
||||
WARNING: These APIs are experimental and may change without notice.
|
||||
"""
|
||||
|
||||
# We create this inline so we only add these capabilities _if_ they're actually used
|
||||
if self._experimental_handlers is None:
|
||||
self._experimental_handlers = ExperimentalHandlers(self, self.request_handlers, self.notification_handlers)
|
||||
return self._experimental_handlers
|
||||
|
||||
def list_prompts(self):
|
||||
def decorator(
|
||||
func: Callable[[], Awaitable[list[types.Prompt]]]
|
||||
|
|
@ -328,8 +347,6 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
mimeType=mime_type or "text/plain",
|
||||
)
|
||||
case bytes() as data: # pragma: no cover
|
||||
import base64
|
||||
|
||||
return types.BlobResourceContents(
|
||||
uri=req.params.uri,
|
||||
blob=base64.b64encode(data).decode(),
|
||||
|
|
@ -483,7 +500,13 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
def decorator(
|
||||
func: Callable[
|
||||
...,
|
||||
Awaitable[UnstructuredContent | StructuredContent | CombinationContent | types.CallToolResult],
|
||||
Awaitable[
|
||||
UnstructuredContent
|
||||
| StructuredContent
|
||||
| CombinationContent
|
||||
| types.CallToolResult
|
||||
| types.CreateTaskResult
|
||||
],
|
||||
],
|
||||
):
|
||||
logger.debug("Registering handler for CallToolRequest")
|
||||
|
|
@ -509,6 +532,9 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
maybe_structured_content: StructuredContent | None
|
||||
if isinstance(results, types.CallToolResult):
|
||||
return types.ServerResult(results)
|
||||
elif isinstance(results, types.CreateTaskResult):
|
||||
# Task-augmented execution returns task info instead of result
|
||||
return types.ServerResult(results)
|
||||
elif isinstance(results, tuple) and len(results) == 2:
|
||||
# tool returned both structured and unstructured content
|
||||
unstructured_content, maybe_structured_content = cast(CombinationContent, results)
|
||||
|
|
@ -627,6 +653,12 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
)
|
||||
)
|
||||
|
||||
# Configure task support for this session if enabled
|
||||
task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
|
||||
if task_support is not None:
|
||||
task_support.configure_session(session)
|
||||
await stack.enter_async_context(task_support.run())
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
async for message in session.incoming_messages:
|
||||
logger.debug("Received message: %s", message)
|
||||
|
|
@ -669,13 +701,14 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
async def _handle_request(
|
||||
self,
|
||||
message: RequestResponder[types.ClientRequest, types.ServerResult],
|
||||
req: Any,
|
||||
req: types.ClientRequestType,
|
||||
session: ServerSession,
|
||||
lifespan_context: LifespanResultT,
|
||||
raise_exceptions: bool,
|
||||
):
|
||||
logger.info("Processing request of type %s", type(req).__name__)
|
||||
if handler := self.request_handlers.get(type(req)): # type: ignore
|
||||
|
||||
if handler := self.request_handlers.get(type(req)):
|
||||
logger.debug("Dispatching request of type %s", type(req).__name__)
|
||||
|
||||
token = None
|
||||
|
|
@ -689,12 +722,24 @@ class Server(Generic[LifespanResultT, RequestT]):
|
|||
|
||||
# Set our global state that can be retrieved via
|
||||
# app.get_request_context()
|
||||
client_capabilities = session.client_params.capabilities if session.client_params else None
|
||||
task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
|
||||
# Get task metadata from request params if present
|
||||
task_metadata = None
|
||||
if hasattr(req, "params") and req.params is not None:
|
||||
task_metadata = getattr(req.params, "task", None)
|
||||
token = request_ctx.set(
|
||||
RequestContext(
|
||||
message.request_id,
|
||||
message.request_meta,
|
||||
session,
|
||||
lifespan_context,
|
||||
Experimental(
|
||||
task_metadata=task_metadata,
|
||||
_client_capabilities=client_capabilities,
|
||||
_session=session,
|
||||
_task_support=task_support,
|
||||
),
|
||||
request=request_data,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in a new issue