add mvp for server side tasks

Author: maxisbey

src/mcp/server/lowlevel/experimental.py

@@ -0,0 +1,137 @@
+"""Experimental handlers for the low-level MCP server.
+
+WARNING: These APIs are experimental and may change without notice.
+"""
+
+import logging
+from collections.abc import Awaitable, Callable
+
+from mcp.server.lowlevel.func_inspection import create_call_wrapper
+from mcp.types import (
+    CancelTaskRequest,
+    CancelTaskResult,
+    GetTaskPayloadRequest,
+    GetTaskPayloadResult,
+    GetTaskRequest,
+    GetTaskResult,
+    ListTasksRequest,
+    ListTasksResult,
+    ServerCapabilities,
+    ServerResult,
+    ServerTasksCapability,
+    ServerTasksRequestsCapability,
+    TasksCancelCapability,
+    TasksListCapability,
+    TasksToolsCapability,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ExperimentalHandlers:
+    """Experimental request/notification handlers.
+
+    WARNING: These APIs are experimental and may change without notice.
+    """
+
+    def __init__(
+        self,
+        request_handlers: dict[type, Callable[..., Awaitable[ServerResult]]],
+        notification_handlers: dict[type, Callable[..., Awaitable[None]]],
+    ):
+        self._request_handlers = request_handlers
+        self._notification_handlers = notification_handlers
+
+    def update_capabilities(self, capabilities: ServerCapabilities) -> None:
+        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 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):
+                result = await wrapper(req)
+                return ServerResult(result)
+
+            self._request_handlers[ListTasksRequest] = handler
+            return func
+
+        return decorator
+
+    def get_task(self):
+        """Register a handler for getting task status.
+
+        WARNING: This API is experimental and may change without notice.
+        """
+
+        def decorator(func: Callable[[GetTaskRequest], Awaitable[GetTaskResult]]):
+            logger.debug("Registering handler for GetTaskRequest")
+            wrapper = create_call_wrapper(func, GetTaskRequest)
+
+            async def handler(req: GetTaskRequest):
+                result = await wrapper(req)
+                return ServerResult(result)
+
+            self._request_handlers[GetTaskRequest] = handler
+            return func
+
+        return decorator
+
+    def get_task_result(self):
+        """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]]):
+            logger.debug("Registering handler for GetTaskPayloadRequest")
+            wrapper = create_call_wrapper(func, GetTaskPayloadRequest)
+
+            async def handler(req: GetTaskPayloadRequest):
+                result = await wrapper(req)
+                return ServerResult(result)
+
+            self._request_handlers[GetTaskPayloadRequest] = handler
+            return func
+
+        return decorator
+
+    def cancel_task(self):
+        """Register a handler for cancelling tasks.
+
+        WARNING: This API is experimental and may change without notice.
+        """
+
+        def decorator(func: Callable[[CancelTaskRequest], Awaitable[CancelTaskResult]]):
+            logger.debug("Registering handler for CancelTaskRequest")
+            wrapper = create_call_wrapper(func, CancelTaskRequest)
+
+            async def handler(req: CancelTaskRequest):
+                result = await wrapper(req)
+                return ServerResult(result)
+
+            self._request_handlers[CancelTaskRequest] = handler
+            return func
+
+        return decorator

src/mcp/server/lowlevel/server.py

@@ -82,11 +82,12 @@ async def main():
 from typing_extensions import TypeVar
 
 import mcp.types as types
+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
 from mcp.server.session import ServerSession
-from mcp.shared.context import RequestContext
+from mcp.shared.context import Experimental, RequestContext
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.session import RequestResponder
@@ -154,6 +155,7 @@ def __init__(
         }
         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(
@@ -219,14 +221,17 @@ def get_capabilities(
         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,
             logging=logging_capability,
             experimental=experimental_capabilities,
             completions=completions_capability,
         )
+        if self._experimental_handlers:
+            self._experimental_handlers.update_capabilities(capabilities)
+        return capabilities
 
     @property
     def request_context(
@@ -235,6 +240,18 @@ def request_context(
         """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.request_handlers, self.notification_handlers)
+        return self._experimental_handlers
+
     def list_prompts(self):
         def decorator(
             func: Callable[[], Awaitable[list[types.Prompt]]]
@@ -666,13 +683,14 @@ async def _handle_message(
     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
@@ -692,6 +710,7 @@ async def _handle_request(
                         message.request_meta,
                         session,
                         lifespan_context,
+                        Experimental(task_metadata=message.request_params.task if message.request_params else None),
                         request=request_data,
                     )
                 )

src/mcp/shared/context.py

@@ -4,17 +4,27 @@
 from typing_extensions import TypeVar
 
 from mcp.shared.session import BaseSession
-from mcp.types import RequestId, RequestParams
+from mcp.types import RequestId, RequestParams, TaskMetadata
 
 SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
 LifespanContextT = TypeVar("LifespanContextT")
 RequestT = TypeVar("RequestT", default=Any)
 
 
+@dataclass
+class Experimental:
+    task_metadata: TaskMetadata | None = None
+
+    @property
+    def is_task(self) -> bool:
+        return self.task_metadata is not None
+
+
 @dataclass
 class RequestContext(Generic[SessionT, LifespanContextT, RequestT]):
     request_id: RequestId
     meta: RequestParams.Meta | None
     session: SessionT
     lifespan_context: LifespanContextT
+    experimental: Experimental = Experimental()
     request: RequestT | None = None

src/mcp/shared/session.py

@@ -81,9 +81,11 @@ def __init__(
         ]""",
         on_complete: Callable[["RequestResponder[ReceiveRequestT, SendResultT]"], Any],
         message_metadata: MessageMetadata = None,
+        request_params: RequestParams | None = None,
     ) -> None:
         self.request_id = request_id
         self.request_meta = request_meta
+        self.request_params = request_params
         self.request = request
         self.message_metadata = message_metadata
         self._session = session
@@ -353,6 +355,7 @@ async def _receive_loop(self) -> None:
                                 session=self,
                                 on_complete=lambda r: self._in_flight.pop(r.request_id, None),
                                 message_metadata=message.metadata,
+                                request_params=validated_request.root.params,
                             )
                             self._in_flight[responder.request_id] = responder
                             await self._received_request(responder)