tasks additions

Author: maxisbey

src/mcp/server/lowlevel/server.py

@@ -497,7 +497,13 @@ def call_tool(self, *, validate_input: bool = True):
         def decorator(
             func: Callable[
                 ...,
-                Awaitable[UnstructuredContent | StructuredContent | CombinationContent | types.CallToolResult],
+                Awaitable[
+                    UnstructuredContent
+                    | StructuredContent
+                    | CombinationContent
+                    | types.CallToolResult
+                    | types.CreateTaskResult
+                ],
             ],
         ):
             logger.debug("Registering handler for CallToolRequest")
@@ -523,6 +529,9 @@ async def handler(req: types.CallToolRequest):
                     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)

src/mcp/shared/context.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Generic
 
 from typing_extensions import TypeVar
@@ -26,5 +26,5 @@ class RequestContext(Generic[SessionT, LifespanContextT, RequestT]):
     meta: RequestParams.Meta | None
     session: SessionT
     lifespan_context: LifespanContextT
-    experimental: Experimental = Experimental()
+    experimental: Experimental = field(default_factory=Experimental)
     request: RequestT | None = None

src/mcp/shared/experimental/__init__.py

@@ -0,0 +1,8 @@
+"""Experimental MCP features.
+
+WARNING: These APIs are experimental and may change without notice.
+"""
+
+from mcp.shared.experimental import tasks
+
+__all__ = ["tasks"]

src/mcp/shared/experimental/tasks/__init__.py

@@ -0,0 +1,38 @@
+"""
+Experimental task management for MCP.
+
+This module provides:
+- TaskStore: Abstract interface for task state storage
+- TaskContext: Context object for task work to interact with state/notifications
+- InMemoryTaskStore: Reference implementation for testing/development
+- Helper functions: run_task, is_terminal, create_task_state, generate_task_id
+
+Architecture:
+- TaskStore is pure storage - it doesn't know about execution
+- TaskContext wraps store + session, providing a clean API for task work
+- run_task is optional convenience for spawning in-process tasks
+
+WARNING: These APIs are experimental and may change without notice.
+"""
+
+from mcp.shared.experimental.tasks.context import TaskContext
+from mcp.shared.experimental.tasks.helpers import (
+    create_task_state,
+    generate_task_id,
+    is_terminal,
+    run_task,
+    task_execution,
+)
+from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
+from mcp.shared.experimental.tasks.store import TaskStore
+
+__all__ = [
+    "TaskStore",
+    "TaskContext",
+    "InMemoryTaskStore",
+    "run_task",
+    "task_execution",
+    "is_terminal",
+    "create_task_state",
+    "generate_task_id",
+]

src/mcp/shared/experimental/tasks/context.py

@@ -0,0 +1,140 @@
+"""
+TaskContext - Context for task work to interact with state and notifications.
+"""
+
+from typing import TYPE_CHECKING
+
+from mcp.shared.experimental.tasks.store import TaskStore
+from mcp.types import (
+    Result,
+    ServerNotification,
+    Task,
+    TaskStatusNotification,
+    TaskStatusNotificationParams,
+)
+
+if TYPE_CHECKING:
+    from mcp.server.session import ServerSession
+
+
+class TaskContext:
+    """
+    Context provided to task work for state management and notifications.
+
+    This wraps a TaskStore and optional session, providing a clean API
+    for task work to update status, complete, fail, and send notifications.
+
+    Example:
+        async def my_task_work(ctx: TaskContext) -> CallToolResult:
+            await ctx.update_status("Starting processing...")
+
+            for i, item in enumerate(items):
+                await ctx.update_status(f"Processing item {i+1}/{len(items)}")
+                if ctx.is_cancelled:
+                    return CallToolResult(content=[TextContent(type="text", text="Cancelled")])
+                process(item)
+
+            return CallToolResult(content=[TextContent(type="text", text="Done!")])
+    """
+
+    def __init__(
+        self,
+        task: Task,
+        store: TaskStore,
+        session: "ServerSession | None" = None,
+    ):
+        self._task = task
+        self._store = store
+        self._session = session
+        self._cancelled = False
+
+    @property
+    def task_id(self) -> str:
+        """The task identifier."""
+        return self._task.taskId
+
+    @property
+    def task(self) -> Task:
+        """The current task state."""
+        return self._task
+
+    @property
+    def is_cancelled(self) -> bool:
+        """Whether cancellation has been requested."""
+        return self._cancelled
+
+    def request_cancellation(self) -> None:
+        """
+        Request cancellation of this task.
+
+        This sets is_cancelled=True. Task work should check this
+        periodically and exit gracefully if set.
+        """
+        self._cancelled = True
+
+    async def update_status(self, message: str, *, notify: bool = True) -> None:
+        """
+        Update the task's status message.
+
+        Args:
+            message: The new status message
+            notify: Whether to send a notification to the client
+        """
+        self._task = await self._store.update_task(
+            self.task_id,
+            status_message=message,
+        )
+        if notify:
+            await self._send_notification()
+
+    async def complete(self, result: Result, *, notify: bool = True) -> None:
+        """
+        Mark the task as completed with the given result.
+
+        Args:
+            result: The task result
+            notify: Whether to send a notification to the client
+        """
+        await self._store.store_result(self.task_id, result)
+        self._task = await self._store.update_task(
+            self.task_id,
+            status="completed",
+        )
+        if notify:
+            await self._send_notification()
+
+    async def fail(self, error: str, *, notify: bool = True) -> None:
+        """
+        Mark the task as failed with an error message.
+
+        Args:
+            error: The error message
+            notify: Whether to send a notification to the client
+        """
+        self._task = await self._store.update_task(
+            self.task_id,
+            status="failed",
+            status_message=error,
+        )
+        if notify:
+            await self._send_notification()
+
+    async def _send_notification(self) -> None:
+        """Send a task status notification to the client."""
+        if self._session is None:
+            return
+
+        await self._session.send_notification(
+            ServerNotification(
+                TaskStatusNotification(
+                    params=TaskStatusNotificationParams(
+                        taskId=self._task.taskId,
+                        status=self._task.status,
+                        statusMessage=self._task.statusMessage,
+                        createdAt=self._task.createdAt,
+                        ttl=self._task.ttl,
+                        pollInterval=self._task.pollInterval,
+                    )
+                )
+            )
+        )