Refactor tasks capability checking into isolated module

maxisbey · maxisbey · commit 8cd276569a6c · 2025-11-28T14:55:13.000Z
Move all task-related capability checking logic into
mcp/shared/experimental/tasks/capabilities.py to keep tasks code
isolated from core session code.

Changes:
- Create capabilities.py with check_tasks_capability() and require_* helpers
- Update ServerSession to import and use the shared function
- Update ServerTaskContext to use require_* helpers instead of inline checks
- Add missing capability checks to ExperimentalServerSessionFeatures

This improves code organization and fixes a bug where
session.experimental.elicit_as_task() wasn't checking capabilities.
diff --git a/src/mcp/server/experimental/session_features.py b/src/mcp/server/experimental/session_features.py
@@ -11,6 +11,10 @@
 from typing import TYPE_CHECKING, Any, TypeVar
 
 import mcp.types as types
+from mcp.shared.experimental.tasks.capabilities import (
+    require_task_augmented_elicitation,
+    require_task_augmented_sampling,
+)
 from mcp.shared.experimental.tasks.polling import poll_until_terminal
 
 if TYPE_CHECKING:
@@ -113,7 +117,13 @@ async def elicit_as_task(
 
         Returns:
             The client's elicitation response
+
+        Raises:
+            McpError: If client doesn't support task-augmented elicitation
         """
+        client_caps = self._session.client_params.capabilities if self._session.client_params else None
+        require_task_augmented_elicitation(client_caps)
+
         create_result = await self._session.send_request(
             types.ServerRequest(
                 types.ElicitRequest(
@@ -166,7 +176,13 @@ async def create_message_as_task(
 
         Returns:
             The sampling result from the client
+
+        Raises:
+            McpError: If client doesn't support task-augmented sampling
         """
+        client_caps = self._session.client_params.capabilities if self._session.client_params else None
+        require_task_augmented_sampling(client_caps)
+
         create_result = await self._session.send_request(
             types.ServerRequest(
                 types.CreateMessageRequest(
diff --git a/src/mcp/server/experimental/task_context.py b/src/mcp/server/experimental/task_context.py
@@ -14,6 +14,10 @@
 from mcp.server.experimental.task_result_handler import TaskResultHandler
 from mcp.server.session import ServerSession
 from mcp.shared.exceptions import McpError
+from mcp.shared.experimental.tasks.capabilities import (
+    require_task_augmented_elicitation,
+    require_task_augmented_sampling,
+)
 from mcp.shared.experimental.tasks.context import TaskContext
 from mcp.shared.experimental.tasks.message_queue import QueuedMessage, TaskMessageQueue
 from mcp.shared.experimental.tasks.resolver import Resolver
@@ -23,8 +27,6 @@
     TASK_STATUS_INPUT_REQUIRED,
     TASK_STATUS_WORKING,
     ClientCapabilities,
-    ClientTasksCapability,
-    ClientTasksRequestsCapability,
     CreateMessageResult,
     CreateTaskResult,
     ElicitationCapability,
@@ -40,10 +42,6 @@
     ServerNotification,
     Task,
     TaskMetadata,
-    TasksCreateElicitationCapability,
-    TasksCreateMessageCapability,
-    TasksElicitationCapability,
-    TasksSamplingCapability,
     TaskStatusNotification,
     TaskStatusNotificationParams,
 )
@@ -198,40 +196,6 @@ def _check_sampling_capability(self) -> None:
                 )
             )
 
-    def _check_task_augmented_elicitation_capability(self) -> None:
-        """Check if the client supports task-augmented elicitation."""
-        capability = ClientCapabilities(
-            tasks=ClientTasksCapability(
-                requests=ClientTasksRequestsCapability(
-                    elicitation=TasksElicitationCapability(create=TasksCreateElicitationCapability())
-                )
-            )
-        )
-        if not self._session.check_client_capability(capability):
-            raise McpError(
-                ErrorData(
-                    code=INVALID_REQUEST,
-                    message="Client does not support task-augmented elicitation capability",
-                )
-            )
-
-    def _check_task_augmented_sampling_capability(self) -> None:
-        """Check if the client supports task-augmented sampling."""
-        capability = ClientCapabilities(
-            tasks=ClientTasksCapability(
-                requests=ClientTasksRequestsCapability(
-                    sampling=TasksSamplingCapability(createMessage=TasksCreateMessageCapability())
-                )
-            )
-        )
-        if not self._session.check_client_capability(capability):
-            raise McpError(
-                ErrorData(
-                    code=INVALID_REQUEST,
-                    message="Client does not support task-augmented sampling capability",
-                )
-            )
-
     async def elicit(
         self,
         message: str,
@@ -412,7 +376,8 @@ async def elicit_as_task(
             McpError: If client doesn't support task-augmented elicitation
             RuntimeError: If handler is not configured
         """
-        self._check_task_augmented_elicitation_capability()
+        client_caps = self._session.client_params.capabilities if self._session.client_params else None
+        require_task_augmented_elicitation(client_caps)
 
         if self._handler is None:
             raise RuntimeError("handler is required for elicit_as_task()")
@@ -504,7 +469,8 @@ async def create_message_as_task(
             McpError: If client doesn't support task-augmented sampling
             RuntimeError: If handler is not configured
         """
-        self._check_task_augmented_sampling_capability()
+        client_caps = self._session.client_params.capabilities if self._session.client_params else None
+        require_task_augmented_sampling(client_caps)
 
         if self._handler is None:
             raise RuntimeError("handler is required for create_message_as_task()")
diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py
@@ -49,6 +49,7 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures
 from mcp.server.models import InitializationOptions
 from mcp.shared.exceptions import McpError
+from mcp.shared.experimental.tasks.capabilities import check_tasks_capability
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.response_router import ResponseRouter
 from mcp.shared.session import (
@@ -116,32 +117,6 @@ def experimental(self) -> ExperimentalServerSessionFeatures:
             self._experimental_features = ExperimentalServerSessionFeatures(self)
         return self._experimental_features
 
-    def _check_tasks_capability(
-        self,
-        required: types.ClientTasksCapability,
-        client: types.ClientTasksCapability,
-    ) -> bool:  # pragma: no cover
-        """Check if client's tasks capability matches the required capability."""
-        if required.requests is None:
-            return True
-        if client.requests is None:
-            return False
-        # Check elicitation.create
-        if required.requests.elicitation is not None:
-            if client.requests.elicitation is None:
-                return False
-            if required.requests.elicitation.create is not None:
-                if client.requests.elicitation.create is None:
-                    return False
-        # Check sampling.createMessage
-        if required.requests.sampling is not None:
-            if client.requests.sampling is None:
-                return False
-            if required.requests.sampling.createMessage is not None:
-                if client.requests.sampling.createMessage is None:
-                    return False
-        return True
-
     def check_client_capability(self, capability: types.ClientCapabilities) -> bool:  # pragma: no cover
         """Check if the client supports a specific capability."""
         if self._client_params is None:
@@ -176,7 +151,7 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool:
         if capability.tasks is not None:
             if client_caps.tasks is None:
                 return False
-            if not self._check_tasks_capability(capability.tasks, client_caps.tasks):
+            if not check_tasks_capability(capability.tasks, client_caps.tasks):
                 return False
 
         return True
diff --git a/src/mcp/shared/experimental/tasks/capabilities.py b/src/mcp/shared/experimental/tasks/capabilities.py
@@ -0,0 +1,115 @@
+"""
+Tasks capability checking utilities.
+
+This module provides functions for checking and requiring task-related
+capabilities. All tasks capability logic is centralized here to keep
+the main session code clean.
+
+WARNING: These APIs are experimental and may change without notice.
+"""
+
+from mcp.shared.exceptions import McpError
+from mcp.types import (
+    INVALID_REQUEST,
+    ClientCapabilities,
+    ClientTasksCapability,
+    ErrorData,
+)
+
+
+def check_tasks_capability(
+    required: ClientTasksCapability,
+    client: ClientTasksCapability,
+) -> bool:
+    """
+    Check if client's tasks capability matches the required capability.
+
+    Args:
+        required: The capability being checked for
+        client: The client's declared capabilities
+
+    Returns:
+        True if client has the required capability, False otherwise
+    """
+    if required.requests is None:
+        return True
+    if client.requests is None:
+        return False
+
+    # Check elicitation.create
+    if required.requests.elicitation is not None:
+        if client.requests.elicitation is None:
+            return False
+        if required.requests.elicitation.create is not None:
+            if client.requests.elicitation.create is None:
+                return False
+
+    # Check sampling.createMessage
+    if required.requests.sampling is not None:
+        if client.requests.sampling is None:
+            return False
+        if required.requests.sampling.createMessage is not None:
+            if client.requests.sampling.createMessage is None:
+                return False
+
+    return True
+
+
+def has_task_augmented_elicitation(caps: ClientCapabilities) -> bool:
+    """Check if capabilities include task-augmented elicitation support."""
+    if caps.tasks is None:
+        return False
+    if caps.tasks.requests is None:
+        return False
+    if caps.tasks.requests.elicitation is None:
+        return False
+    return caps.tasks.requests.elicitation.create is not None
+
+
+def has_task_augmented_sampling(caps: ClientCapabilities) -> bool:
+    """Check if capabilities include task-augmented sampling support."""
+    if caps.tasks is None:
+        return False
+    if caps.tasks.requests is None:
+        return False
+    if caps.tasks.requests.sampling is None:
+        return False
+    return caps.tasks.requests.sampling.createMessage is not None
+
+
+def require_task_augmented_elicitation(client_caps: ClientCapabilities | None) -> None:
+    """
+    Raise McpError if client doesn't support task-augmented elicitation.
+
+    Args:
+        client_caps: The client's declared capabilities, or None if not initialized
+
+    Raises:
+        McpError: If client doesn't support task-augmented elicitation
+    """
+    if client_caps is None or not has_task_augmented_elicitation(client_caps):
+        raise McpError(
+            ErrorData(
+                code=INVALID_REQUEST,
+                message="Client does not support task-augmented elicitation",
+            )
+        )
+
+
+def require_task_augmented_sampling(client_caps: ClientCapabilities | None) -> None:
+    """
+    Raise McpError if client doesn't support task-augmented sampling.
+
+    Args:
+        client_caps: The client's declared capabilities, or None if not initialized
+
+    Raises:
+        McpError: If client doesn't support task-augmented sampling
+    """
+    if client_caps is None or not has_task_augmented_sampling(client_caps):
+        raise McpError(
+            ErrorData(
+                code=INVALID_REQUEST,
+                message="Client does not support task-augmented sampling",
+            )
+        )
