SEP-1686: Tasks (#1645)

src/mcp/server/experimental/task_result_handler.py

@@ -0,0 +1,235 @@
+"""
+TaskResultHandler - Integrated handler for tasks/result endpoint.
+
+This implements the dequeue-send-wait pattern from the MCP Tasks spec:
+1. Dequeue all pending messages for the task
+2. Send them to the client via transport with relatedRequestId routing
+3. Wait if task is not in terminal state
+4. Return final result when task completes
+
+This is the core of the task message queue pattern.
+"""
+
+import logging
+from typing import Any
+
+import anyio
+
+from mcp.server.session import ServerSession
+from mcp.shared.exceptions import McpError
+from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY, is_terminal
+from mcp.shared.experimental.tasks.message_queue import TaskMessageQueue
+from mcp.shared.experimental.tasks.resolver import Resolver
+from mcp.shared.experimental.tasks.store import TaskStore
+from mcp.shared.message import ServerMessageMetadata, SessionMessage
+from mcp.types import (
+    INVALID_PARAMS,
+    ErrorData,
+    GetTaskPayloadRequest,
+    GetTaskPayloadResult,
+    JSONRPCMessage,
+    RelatedTaskMetadata,
+    RequestId,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class TaskResultHandler:
+    """
+    Handler for tasks/result that implements the message queue pattern.
+
+    This handler:
+    1. Dequeues pending messages (elicitations, notifications) for the task
+    2. Sends them to the client via the response stream
+    3. Waits for responses and resolves them back to callers
+    4. Blocks until task reaches terminal state
+    5. Returns the final result
+
+    Usage:
+        # Create handler with store and queue
+        handler = TaskResultHandler(task_store, message_queue)
+
+        # Register it with the server
+        @server.experimental.get_task_result()
+        async def handle_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
+            ctx = server.request_context
+            return await handler.handle(req, ctx.session, ctx.request_id)
+
+        # Or use the convenience method
+        handler.register(server)
+    """
+
+    def __init__(
+        self,
+        store: TaskStore,
+        queue: TaskMessageQueue,
+    ):
+        self._store = store
+        self._queue = queue
+        # Map from internal request ID to resolver for routing responses
+        self._pending_requests: dict[RequestId, Resolver[dict[str, Any]]] = {}
+
+    async def send_message(
+        self,
+        session: ServerSession,
+        message: SessionMessage,
+    ) -> None:
+        """
+        Send a message via the session.
+
+        This is a helper for delivering queued task messages.
+        """
+        await session.send_message(message)
+
+    async def handle(
+        self,
+        request: GetTaskPayloadRequest,
+        session: ServerSession,
+        request_id: RequestId,
+    ) -> GetTaskPayloadResult:
+        """
+        Handle a tasks/result request.
+
+        This implements the dequeue-send-wait loop:
+        1. Dequeue all pending messages
+        2. Send each via transport with relatedRequestId = this request's ID
+        3. If task not terminal, wait for status change
+        4. Loop until task is terminal
+        5. Return final result
+
+        Args:
+            request: The GetTaskPayloadRequest
+            session: The server session for sending messages
+            request_id: The request ID for relatedRequestId routing
+
+        Returns:
+            GetTaskPayloadResult with the task's final payload
+        """
+        task_id = request.params.taskId
+
+        while True:
+            task = await self._store.get_task(task_id)
+            if task is None:
+                raise McpError(
+                    ErrorData(
+                        code=INVALID_PARAMS,
+                        message=f"Task not found: {task_id}",
+                    )
+                )
+
+            await self._deliver_queued_messages(task_id, session, request_id)
+
+            # If task is terminal, return result
+            if is_terminal(task.status):
+                result = await self._store.get_result(task_id)
+                # GetTaskPayloadResult is a Result with extra="allow"
+                # The stored result contains the actual payload data
+                # Per spec: tasks/result MUST include _meta with related-task metadata
+                related_task = RelatedTaskMetadata(taskId=task_id)
+                related_task_meta: dict[str, Any] = {RELATED_TASK_METADATA_KEY: related_task.model_dump(by_alias=True)}
+                if result is not None:
+                    result_data = result.model_dump(by_alias=True)
+                    existing_meta: dict[str, Any] = result_data.get("_meta") or {}
+                    result_data["_meta"] = {**existing_meta, **related_task_meta}
+                    return GetTaskPayloadResult.model_validate(result_data)
+                return GetTaskPayloadResult.model_validate({"_meta": related_task_meta})
+
+            # Wait for task update (status change or new messages)
+            await self._wait_for_task_update(task_id)
+
+    async def _deliver_queued_messages(
+        self,
+        task_id: str,
+        session: ServerSession,
+        request_id: RequestId,
+    ) -> None:
+        """
+        Dequeue and send all pending messages for a task.
+
+        Each message is sent via the session's write stream with
+        relatedRequestId set so responses route back to this stream.
+        """
+        while True:
+            message = await self._queue.dequeue(task_id)
+            if message is None:
+                break
+
+            # If this is a request (not notification), wait for response
+            if message.type == "request" and message.resolver is not None:
+                # Store the resolver so we can route the response back
+                original_id = message.original_request_id
+                if original_id is not None:
+                    self._pending_requests[original_id] = message.resolver
+
+            logger.debug("Delivering queued message for task %s: %s", task_id, message.type)
+
+            # Send the message with relatedRequestId for routing
+            session_message = SessionMessage(
+                message=JSONRPCMessage(message.message),
+                metadata=ServerMessageMetadata(related_request_id=request_id),
+            )
+            await self.send_message(session, session_message)
+
+    async def _wait_for_task_update(self, task_id: str) -> None:
+        """
+        Wait for task to be updated (status change or new message).
+
+        Races between store update and queue message - first one wins.
+        """
+        async with anyio.create_task_group() as tg:
+
+            async def wait_for_store() -> None:
+                try:
+                    await self._store.wait_for_update(task_id)
+                except Exception:
+                    pass
+                finally:
+                    tg.cancel_scope.cancel()
+
+            async def wait_for_queue() -> None:
+                try:
+                    await self._queue.wait_for_message(task_id)
+                except Exception:
+                    pass
+                finally:
+                    tg.cancel_scope.cancel()
+
+            tg.start_soon(wait_for_store)
+            tg.start_soon(wait_for_queue)
+
+    def route_response(self, request_id: RequestId, response: dict[str, Any]) -> bool:
+        """
+        Route a response back to the waiting resolver.
+
+        This is called when a response arrives for a queued request.
+
+        Args:
+            request_id: The request ID from the response
+            response: The response data
+
+        Returns:
+            True if response was routed, False if no pending request
+        """
+        resolver = self._pending_requests.pop(request_id, None)
+        if resolver is not None and not resolver.done():
+            resolver.set_result(response)
+            return True
+        return False
+
+    def route_error(self, request_id: RequestId, error: ErrorData) -> bool:
+        """
+        Route an error back to the waiting resolver.
+
+        Args:
+            request_id: The request ID from the error response
+            error: The error data
+
+        Returns:
+            True if error was routed, False if no pending request
+        """
+        resolver = self._pending_requests.pop(request_id, None)
+        if resolver is not None and not resolver.done():
+            resolver.set_exception(McpError(error))
+            return True
+        return False