Implement capability gating for tasks

Author: LucaButBoring

src/mcp/client/session.py

@@ -15,6 +15,7 @@
 from mcp.shared.context import RequestContext
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
+from mcp.shared.task import TaskStore
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 
 if TYPE_CHECKING:
@@ -126,13 +127,15 @@ def __init__(
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
+        task_store: TaskStore | None = None,
     ) -> None:
         super().__init__(
             read_stream,
             write_stream,
             types.ServerRequest,
             types.ServerNotification,
             read_timeout_seconds=read_timeout_seconds,
+            task_store=task_store,
         )
         self._client_info = client_info or DEFAULT_CLIENT_INFO
         self._sampling_callback = sampling_callback or _default_sampling_callback
@@ -156,6 +159,18 @@ async def initialize(self) -> types.InitializeResult:
             else None
         )
 
+        # Build tasks capability - only if task store is configured
+        tasks = None
+        if self._task_store is not None:
+            tasks = types.ClientTasksCapability(
+                requests=types.ClientTasksRequestsCapability(
+                    sampling=types.TaskSamplingCapability(createMessage=True),
+                    elicitation=types.TaskElicitationCapability(create=True),
+                    roots=types.TaskRootsCapability(list=True),
+                    tasks=types.TasksOperationsCapability(get=True, list=True, result=True, delete=True),
+                )
+            )
+
         result = await self.send_request(
             types.ClientRequest(
                 types.InitializeRequest(
@@ -166,6 +181,7 @@ async def initialize(self) -> types.InitializeResult:
                             elicitation=elicitation,
                             experimental=None,
                             roots=roots,
+                            tasks=tasks,
                         ),
                         clientInfo=self._client_info,
                     ),

src/mcp/server/lowlevel/server.py

@@ -112,10 +112,12 @@ def __init__(
         prompts_changed: bool = False,
         resources_changed: bool = False,
         tools_changed: bool = False,
+        tasks_changed: bool = False,
     ):
         self.prompts_changed = prompts_changed
         self.resources_changed = resources_changed
         self.tools_changed = tools_changed
+        self.tasks_changed = tasks_changed
 
 
 @asynccontextmanager
@@ -199,6 +201,7 @@ def get_capabilities(
         tools_capability = None
         logging_capability = None
         completions_capability = None
+        tasks_capability = None
 
         # Set prompt capabilities if handler exists
         if types.ListPromptsRequest in self.request_handlers:
@@ -222,13 +225,59 @@ def get_capabilities(
         if types.CompleteRequest in self.request_handlers:
             completions_capability = types.CompletionsCapability()
 
+        # Set tasks capabilities if task store is configured
+        if self.task_store is not None:
+            # Build nested request capabilities based on available handlers
+            tools_req_cap = None
+            resources_req_cap = None
+            prompts_req_cap = None
+            tasks_ops_cap = None
+
+            # Check for tool capabilities
+            has_call_tool = types.CallToolRequest in self.request_handlers
+            has_list_tools = types.ListToolsRequest in self.request_handlers
+            if has_call_tool or has_list_tools:
+                tools_req_cap = types.TaskToolsCapability(
+                    call=True if has_call_tool else None, list=True if has_list_tools else None
+                )
+
+            # Check for resource capabilities
+            has_read_resource = types.ReadResourceRequest in self.request_handlers
+            has_list_resources = types.ListResourcesRequest in self.request_handlers
+            if has_read_resource or has_list_resources:
+                resources_req_cap = types.TaskResourcesCapability(
+                    read=True if has_read_resource else None, list=True if has_list_resources else None
+                )
+
+            # Check for prompt capabilities
+            has_get_prompt = types.GetPromptRequest in self.request_handlers
+            has_list_prompts = types.ListPromptsRequest in self.request_handlers
+            if has_get_prompt or has_list_prompts:
+                prompts_req_cap = types.TaskPromptsCapability(
+                    get=True if has_get_prompt else None, list=True if has_list_prompts else None
+                )
+
+            # Task operations are always available if task_store is configured
+            tasks_ops_cap = types.TasksOperationsCapability(get=True, list=True, result=True, delete=True)
+
+            # Build the nested tasks capability
+            tasks_capability = types.ServerTasksCapability(
+                requests=types.ServerTasksRequestsCapability(
+                    tools=tools_req_cap,
+                    resources=resources_req_cap,
+                    prompts=prompts_req_cap,
+                    tasks=tasks_ops_cap,
+                )
+            )
+
         return types.ServerCapabilities(
             prompts=prompts_capability,
             resources=resources_capability,
             tools=tools_capability,
             logging=logging_capability,
             experimental=experimental_capabilities,
             completions=completions_capability,
+            tasks=tasks_capability,
         )
 
     @property

src/mcp/server/session.py

@@ -107,6 +107,52 @@ def __init__(
     def client_params(self) -> types.InitializeRequestParams | None:
         return self._client_params
 
+    def _check_tasks_capability(
+        self, required: types.ClientTasksCapability, client: types.ClientTasksCapability
+    ) -> bool:
+        """Check if client supports required tasks capabilities."""
+        if required.requests is None:
+            return True
+        if client.requests is None:
+            return False
+
+        req_cap = required.requests
+        client_req_cap = client.requests
+
+        # Check sampling requests
+        if req_cap.sampling is not None and (
+            client_req_cap.sampling is None
+            or (req_cap.sampling.createMessage and not client_req_cap.sampling.createMessage)
+        ):
+            return False
+
+        # Check elicitation requests
+        if req_cap.elicitation is not None and (
+            client_req_cap.elicitation is None or (req_cap.elicitation.create and not client_req_cap.elicitation.create)
+        ):
+            return False
+
+        # Check roots requests
+        if req_cap.roots is not None and (
+            client_req_cap.roots is None or (req_cap.roots.list and not client_req_cap.roots.list)
+        ):
+            return False
+
+        # Check tasks operations
+        if req_cap.tasks is not None:
+            if client_req_cap.tasks is None:
+                return False
+            tasks_checks = [
+                not req_cap.tasks.get or client_req_cap.tasks.get,
+                not req_cap.tasks.list or client_req_cap.tasks.list,
+                not req_cap.tasks.result or client_req_cap.tasks.result,
+                not req_cap.tasks.delete or client_req_cap.tasks.delete,
+            ]
+            if not all(tasks_checks):
+                return False
+
+        return True
+
     def check_client_capability(self, capability: types.ClientCapabilities) -> bool:
         """Check if the client supports a specific capability."""
         if self._client_params is None:
@@ -138,6 +184,12 @@ def check_client_capability(self, capability: types.ClientCapabilities) -> bool:
                 if exp_key not in client_caps.experimental or client_caps.experimental[exp_key] != exp_value:
                     return False
 
+        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):
+                return False
+
         return True
 
     async def _receive_loop(self) -> None:
@@ -193,8 +245,17 @@ async def _received_request(  # noqa: PLR0912
                 # Ping requests are allowed at any time
                 pass
             case types.GetTaskRequest(params=params):
+                # Check if client has announced tasks capability
+                if self._client_params is None or self._client_params.capabilities.tasks is None:
+                    with responder:
+                        await responder.respond(
+                            types.ErrorData(
+                                code=types.INVALID_REQUEST,
+                                message="Client has not announced tasks capability",
+                            )
+                        )
                 # Handle get task requests if task store is available
-                if self._task_store:
+                elif self._task_store:
                     task = await self._task_store.get_task(params.taskId)
                     if task is None:
                         with responder:
@@ -220,8 +281,17 @@ async def _received_request(  # noqa: PLR0912
                             types.ErrorData(code=types.INVALID_REQUEST, message="Task store not configured")
                         )
             case types.GetTaskPayloadRequest(params=params):
+                # Check if client has announced tasks capability
+                if self._client_params is None or self._client_params.capabilities.tasks is None:
+                    with responder:
+                        await responder.respond(
+                            types.ErrorData(
+                                code=types.INVALID_REQUEST,
+                                message="Client has not announced tasks capability",
+                            )
+                        )
                 # Handle get task result requests if task store is available
-                if self._task_store:
+                elif self._task_store:
                     task = await self._task_store.get_task(params.taskId)
                     if task is None:
                         with responder:
@@ -253,8 +323,17 @@ async def _received_request(  # noqa: PLR0912
                             types.ErrorData(code=types.INVALID_REQUEST, message="Task store not configured")
                         )
             case types.ListTasksRequest(params=params):
+                # Check if client has announced tasks capability
+                if self._client_params is None or self._client_params.capabilities.tasks is None:
+                    with responder:
+                        await responder.respond(
+                            types.ErrorData(
+                                code=types.INVALID_REQUEST,
+                                message="Client has not announced tasks capability",
+                            )
+                        )
                 # Handle list tasks requests if task store is available
-                if self._task_store:
+                elif self._task_store:
                     try:
                         result = await self._task_store.list_tasks(params.cursor if params else None)
                         with responder:
@@ -278,8 +357,17 @@ async def _received_request(  # noqa: PLR0912
                             types.ErrorData(code=types.INVALID_REQUEST, message="Task store not configured")
                         )
             case types.DeleteTaskRequest(params=params):
+                # Check if client has announced tasks capability
+                if self._client_params is None or self._client_params.capabilities.tasks is None:
+                    with responder:
+                        await responder.respond(
+                            types.ErrorData(
+                                code=types.INVALID_REQUEST,
+                                message="Client has not announced tasks capability",
+                            )
+                        )
                 # Handle delete task requests if task store is available
-                if self._task_store:
+                elif self._task_store:
                     try:
                         await self._task_store.delete_task(params.taskId)
                         with responder: