notifications and client side

Author: maxisbey

src/mcp/client/session.py

@@ -49,6 +49,95 @@ async def __call__(
     ) -> None: ...  # pragma: no branch
 
 
+# Experimental: Task handler protocols for server -> client requests
+class GetTaskHandlerFnT(Protocol):
+    """Handler for tasks/get requests from server.
+
+    WARNING: This is experimental and may change without notice.
+    """
+
+    async def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.GetTaskRequestParams,
+    ) -> types.GetTaskResult | types.ErrorData: ...  # pragma: no branch
+
+
+class GetTaskResultHandlerFnT(Protocol):
+    """Handler for tasks/result requests from server.
+
+    WARNING: This is experimental and may change without notice.
+    """
+
+    async def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.GetTaskPayloadRequestParams,
+    ) -> types.GetTaskPayloadResult | types.ErrorData: ...  # pragma: no branch
+
+
+class ListTasksHandlerFnT(Protocol):
+    """Handler for tasks/list requests from server.
+
+    WARNING: This is experimental and may change without notice.
+    """
+
+    async def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.PaginatedRequestParams | None,
+    ) -> types.ListTasksResult | types.ErrorData: ...  # pragma: no branch
+
+
+class CancelTaskHandlerFnT(Protocol):
+    """Handler for tasks/cancel requests from server.
+
+    WARNING: This is experimental and may change without notice.
+    """
+
+    async def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.CancelTaskRequestParams,
+    ) -> types.CancelTaskResult | types.ErrorData: ...  # pragma: no branch
+
+
+class TaskAugmentedSamplingFnT(Protocol):
+    """Handler for task-augmented sampling/createMessage requests from server.
+
+    When server sends a CreateMessageRequest with task field, this callback
+    is invoked. The callback should create a task, spawn background work,
+    and return CreateTaskResult immediately.
+
+    WARNING: This is experimental and may change without notice.
+    """
+
+    async def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.CreateMessageRequestParams,
+        task_metadata: types.TaskMetadata,
+    ) -> types.CreateTaskResult | types.ErrorData: ...  # pragma: no branch
+
+
+class TaskAugmentedElicitationFnT(Protocol):
+    """Handler for task-augmented elicitation/create requests from server.
+
+    When server sends an ElicitRequest with task field, this callback
+    is invoked. The callback should create a task, spawn background work,
+    and return CreateTaskResult immediately.
+
+    WARNING: This is experimental and may change without notice.
+    """
+
+    async def __call__(
+        self,
+        context: RequestContext["ClientSession", Any],
+        params: types.ElicitRequestParams,
+        task_metadata: types.TaskMetadata,
+    ) -> types.CreateTaskResult | types.ErrorData: ...  # pragma: no branch
+
+
 class MessageHandlerFnT(Protocol):
     async def __call__(
         self,
@@ -97,6 +186,69 @@ async def _default_logging_callback(
     pass
 
 
+# Default handlers for experimental task requests (return "not supported" errors)
+async def _default_get_task_handler(
+    context: RequestContext["ClientSession", Any],
+    params: types.GetTaskRequestParams,
+) -> types.GetTaskResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.METHOD_NOT_FOUND,
+        message="tasks/get not supported",
+    )
+
+
+async def _default_get_task_result_handler(
+    context: RequestContext["ClientSession", Any],
+    params: types.GetTaskPayloadRequestParams,
+) -> types.GetTaskPayloadResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.METHOD_NOT_FOUND,
+        message="tasks/result not supported",
+    )
+
+
+async def _default_list_tasks_handler(
+    context: RequestContext["ClientSession", Any],
+    params: types.PaginatedRequestParams | None,
+) -> types.ListTasksResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.METHOD_NOT_FOUND,
+        message="tasks/list not supported",
+    )
+
+
+async def _default_cancel_task_handler(
+    context: RequestContext["ClientSession", Any],
+    params: types.CancelTaskRequestParams,
+) -> types.CancelTaskResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.METHOD_NOT_FOUND,
+        message="tasks/cancel not supported",
+    )
+
+
+async def _default_task_augmented_sampling_callback(
+    context: RequestContext["ClientSession", Any],
+    params: types.CreateMessageRequestParams,
+    task_metadata: types.TaskMetadata,
+) -> types.CreateTaskResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.INVALID_REQUEST,
+        message="Task-augmented sampling not supported",
+    )
+
+
+async def _default_task_augmented_elicitation_callback(
+    context: RequestContext["ClientSession", Any],
+    params: types.ElicitRequestParams,
+    task_metadata: types.TaskMetadata,
+) -> types.CreateTaskResult | types.ErrorData:
+    return types.ErrorData(
+        code=types.INVALID_REQUEST,
+        message="Task-augmented elicitation not supported",
+    )
+
+
 ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(types.ClientResult | types.ErrorData)
 
 
@@ -120,6 +272,14 @@ def __init__(
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
+        tasks_capability: types.ClientTasksCapability | None = None,
+        # Experimental: Task handlers for server -> client requests
+        get_task_handler: GetTaskHandlerFnT | None = None,
+        get_task_result_handler: GetTaskResultHandlerFnT | None = None,
+        list_tasks_handler: ListTasksHandlerFnT | None = None,
+        cancel_task_handler: CancelTaskHandlerFnT | None = None,
+        task_augmented_sampling_callback: TaskAugmentedSamplingFnT | None = None,
+        task_augmented_elicitation_callback: TaskAugmentedElicitationFnT | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -134,9 +294,21 @@ def __init__(
         self._list_roots_callback = list_roots_callback or _default_list_roots_callback
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
+        self._tasks_capability = tasks_capability
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
         self._server_capabilities: types.ServerCapabilities | None = None
         self._experimental: ExperimentalClientFeatures | None = None
+        # Experimental: Task handlers
+        self._get_task_handler = get_task_handler or _default_get_task_handler
+        self._get_task_result_handler = get_task_result_handler or _default_get_task_result_handler
+        self._list_tasks_handler = list_tasks_handler or _default_list_tasks_handler
+        self._cancel_task_handler = cancel_task_handler or _default_cancel_task_handler
+        self._task_augmented_sampling_callback = (
+            task_augmented_sampling_callback or _default_task_augmented_sampling_callback
+        )
+        self._task_augmented_elicitation_callback = (
+            task_augmented_elicitation_callback or _default_task_augmented_elicitation_callback
+        )
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
@@ -162,6 +334,7 @@ async def initialize(self) -> types.InitializeResult:
                             elicitation=elicitation,
                             experimental=None,
                             roots=roots,
+                            tasks=self._tasks_capability,
                         ),
                         clientInfo=self._client_info,
                     ),
@@ -187,7 +360,7 @@ def get_server_capabilities(self) -> types.ServerCapabilities | None:
         return self._server_capabilities
 
     @property
-    def experimental(self) -> "ExperimentalClientFeatures":
+    def experimental(self) -> ExperimentalClientFeatures:
         """Experimental APIs for tasks and other features.
 
         WARNING: These APIs are experimental and may change without notice.
@@ -534,13 +707,21 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
         match responder.request.root:
             case types.CreateMessageRequest(params=params):
                 with responder:
-                    response = await self._sampling_callback(ctx, params)
+                    # Check if this is a task-augmented request
+                    if params.task is not None:
+                        response = await self._task_augmented_sampling_callback(ctx, params, params.task)
+                    else:
+                        response = await self._sampling_callback(ctx, params)
                     client_response = ClientResponse.validate_python(response)
                     await responder.respond(client_response)
 
             case types.ElicitRequest(params=params):
                 with responder:
-                    response = await self._elicitation_callback(ctx, params)
+                    # Check if this is a task-augmented request
+                    if params.task is not None:
+                        response = await self._task_augmented_elicitation_callback(ctx, params, params.task)
+                    else:
+                        response = await self._elicitation_callback(ctx, params)
                     client_response = ClientResponse.validate_python(response)
                     await responder.respond(client_response)
 
@@ -553,7 +734,33 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
             case types.PingRequest():  # pragma: no cover
                 with responder:
                     return await responder.respond(types.ClientResult(root=types.EmptyResult()))
-            case _:
+
+            # Experimental: Task management requests from server
+            case types.GetTaskRequest(params=params):
+                with responder:
+                    response = await self._get_task_handler(ctx, params)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
+            case types.GetTaskPayloadRequest(params=params):
+                with responder:
+                    response = await self._get_task_result_handler(ctx, params)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
+            case types.ListTasksRequest(params=params):
+                with responder:
+                    response = await self._list_tasks_handler(ctx, params)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
+            case types.CancelTaskRequest(params=params):
+                with responder:
+                    response = await self._cancel_task_handler(ctx, params)
+                    client_response = ClientResponse.validate_python(response)
+                    await responder.respond(client_response)
+
+            case _:  # pragma: no cover
                 raise NotImplementedError()
 
     async def _handle_incoming(

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

@@ -5,10 +5,13 @@
 - TaskStore: Abstract interface for task state storage
 - TaskContext: Context object for task work to interact with state/notifications
 - InMemoryTaskStore: Reference implementation for testing/development
+- TaskMessageQueue: FIFO queue for task messages delivered via tasks/result
+- InMemoryTaskMessageQueue: Reference implementation for message queue
 - Helper functions: run_task, is_terminal, create_task_state, generate_task_id
 
 Architecture:
 - TaskStore is pure storage - it doesn't know about execution
+- TaskMessageQueue stores messages to be delivered via tasks/result
 - TaskContext wraps store + session, providing a clean API for task work
 - run_task is optional convenience for spawning in-process tasks
 
@@ -24,15 +27,31 @@
     task_execution,
 )
 from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
+from mcp.shared.experimental.tasks.message_queue import (
+    InMemoryTaskMessageQueue,
+    QueuedMessage,
+    TaskMessageQueue,
+)
+from mcp.shared.experimental.tasks.result_handler import (
+    TaskResultHandler,
+    create_task_result_handler,
+)
 from mcp.shared.experimental.tasks.store import TaskStore
+from mcp.shared.experimental.tasks.task_session import TaskSession
 
 __all__ = [
     "TaskStore",
     "TaskContext",
+    "TaskSession",
+    "TaskResultHandler",
     "InMemoryTaskStore",
+    "TaskMessageQueue",
+    "InMemoryTaskMessageQueue",
+    "QueuedMessage",
     "run_task",
     "task_execution",
     "is_terminal",
     "create_task_state",
     "generate_task_id",
+    "create_task_result_handler",
 ]

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

@@ -8,6 +8,7 @@
 For production, consider implementing TaskStore with a database or distributed cache.
 """
 
+import asyncio
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 
@@ -46,6 +47,7 @@ class InMemoryTaskStore(TaskStore):
     def __init__(self, page_size: int = 10) -> None:
         self._tasks: dict[str, StoredTask] = {}
         self._page_size = page_size
+        self._update_events: dict[str, asyncio.Event] = {}
 
     def _calculate_expiry(self, ttl_ms: int | None) -> datetime | None:
         """Calculate expiry time from TTL in milliseconds."""
@@ -111,8 +113,10 @@ async def update_task(
         if stored is None:
             raise ValueError(f"Task with ID {task_id} not found")
 
-        if status is not None:
+        status_changed = False
+        if status is not None and stored.task.status != status:
             stored.task.status = status
+            status_changed = True
 
         if status_message is not None:
             stored.task.statusMessage = status_message
@@ -121,6 +125,10 @@ async def update_task(
         if status is not None and is_terminal(status) and stored.task.ttl is not None:
             stored.expires_at = self._calculate_expiry(stored.task.ttl)
 
+        # Notify waiters if status changed
+        if status_changed:
+            await self.notify_update(task_id)
+
         return Task(**stored.task.model_dump())
 
     async def store_result(self, task_id: str, result: Result) -> None:
@@ -175,11 +183,31 @@ async def delete_task(self, task_id: str) -> bool:
         del self._tasks[task_id]
         return True
 
+    async def wait_for_update(self, task_id: str) -> None:
+        """Wait until the task status changes."""
+        if task_id not in self._tasks:
+            raise ValueError(f"Task with ID {task_id} not found")
+
+        # Get or create the event for this task
+        if task_id not in self._update_events:
+            self._update_events[task_id] = asyncio.Event()
+
+        event = self._update_events[task_id]
+        # Clear before waiting so we wait for NEW updates
+        event.clear()
+        await event.wait()
+
+    async def notify_update(self, task_id: str) -> None:
+        """Signal that a task has been updated."""
+        if task_id in self._update_events:
+            self._update_events[task_id].set()
+
     # --- Testing/debugging helpers ---
 
     def cleanup(self) -> None:
         """Cleanup all tasks (useful for testing or graceful shutdown)."""
         self._tasks.clear()
+        self._update_events.clear()
 
     def get_all_tasks(self) -> list[Task]:
         """Get all tasks (useful for debugging). Returns copies to prevent modification."""