experimental.py

"""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,
    TasksCallCapability,
    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(call=TasksCallCapability())
        )  # 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