add task helpers to context

Author: maxisbey

examples/clients/simple-task-client/mcp_simple_task_client/main.py

@@ -27,6 +27,8 @@ async def run(url: str) -> None:
 
             # Call the tool as a task
             print("\nCalling tool as a task...")
+
+            # TODO: make helper for this
             result = await session.send_request(
                 ClientRequest(
                     CallToolRequest(

examples/servers/simple-task/mcp_simple_task/server.py

@@ -40,6 +40,7 @@ async def list_tools() -> list[types.Tool]:
             name="long_running_task",
             description="A task that takes a few seconds to complete with status updates",
             inputSchema={"type": "object", "properties": {}},
+            execution=types.ToolExecution(task="always"),
         )
     ]
 
@@ -49,8 +50,8 @@ async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[types.T
     ctx = server.request_context
     app = ctx.lifespan_context
 
-    if not ctx.experimental.is_task:
-        return [types.TextContent(type="text", text="Error: This tool must be called as a task")]
+    # Validate task mode - raises McpError(-32601) if client didn't use task augmentation
+    ctx.experimental.validate_task_mode("always")
 
     # Create the task
     metadata = ctx.experimental.task_metadata

src/mcp/server/lowlevel/server.py

@@ -713,13 +713,17 @@ async def _handle_request(
 
                 # Set our global state that can be retrieved via
                 # app.get_request_context()
+                client_capabilities = session.client_params.capabilities if session.client_params else None
                 token = request_ctx.set(
                     RequestContext(
                         message.request_id,
                         message.request_meta,
                         session,
                         lifespan_context,
-                        Experimental(task_metadata=message.request_params.task if message.request_params else None),
+                        Experimental(
+                            task_metadata=message.request_params.task if message.request_params else None,
+                            _client_capabilities=client_capabilities,
+                        ),
                         request=request_data,
                     )
                 )

src/mcp/shared/context.py

@@ -3,8 +3,18 @@
 
 from typing_extensions import TypeVar
 
+from mcp import McpError
 from mcp.shared.session import BaseSession
-from mcp.types import RequestId, RequestParams, TaskMetadata
+from mcp.types import (
+    METHOD_NOT_FOUND,
+    ClientCapabilities,
+    ErrorData,
+    RequestId,
+    RequestParams,
+    TaskExecutionMode,
+    TaskMetadata,
+    Tool,
+)
 
 SessionT = TypeVar("SessionT", bound=BaseSession[Any, Any, Any, Any, Any])
 LifespanContextT = TypeVar("LifespanContextT")
@@ -13,12 +23,111 @@
 
 @dataclass
 class Experimental:
+    """
+    Experimental features context for task-augmented requests.
+
+    Provides helpers for validating task execution compatibility.
+    """
+
     task_metadata: TaskMetadata | None = None
+    _client_capabilities: ClientCapabilities | None = field(default=None, repr=False)
 
     @property
     def is_task(self) -> bool:
+        """Check if this request is task-augmented."""
         return self.task_metadata is not None
 
+    @property
+    def client_supports_tasks(self) -> bool:
+        """Check if the client declared task support."""
+        if self._client_capabilities is None:
+            return False
+        return self._client_capabilities.tasks is not None
+
+    def validate_task_mode(
+        self,
+        tool_task_mode: TaskExecutionMode | None,
+        *,
+        raise_error: bool = True,
+    ) -> ErrorData | None:
+        """
+        Validate that the request is compatible with the tool's task execution mode.
+
+        Per MCP spec:
+        - "always": Clients MUST invoke as task. Server returns -32601 if not.
+        - "never" (or None): Clients MUST NOT invoke as task. Server returns -32601 if they do.
+        - "optional": Either is acceptable.
+
+        Args:
+            tool_task_mode: The tool's execution.task value ("never", "optional", "always", or None)
+            raise_error: If True, raises McpError on validation failure. If False, returns ErrorData.
+
+        Returns:
+            None if valid, ErrorData if invalid and raise_error=False
+
+        Raises:
+            McpError: If invalid and raise_error=True
+        """
+
+        mode = tool_task_mode or "never"
+
+        error: ErrorData | None = None
+
+        if mode == "always" and not self.is_task:
+            error = ErrorData(
+                code=METHOD_NOT_FOUND,
+                message="This tool requires task-augmented invocation",
+            )
+        elif mode == "never" and self.is_task:
+            error = ErrorData(
+                code=METHOD_NOT_FOUND,
+                message="This tool does not support task-augmented invocation",
+            )
+
+        if error is not None and raise_error:
+            raise McpError(error)
+
+        return error
+
+    def validate_for_tool(
+        self,
+        tool: Tool,
+        *,
+        raise_error: bool = True,
+    ) -> ErrorData | None:
+        """
+        Validate that the request is compatible with the given tool.
+
+        Convenience wrapper around validate_task_mode that extracts the mode from a Tool.
+
+        Args:
+            tool: The Tool definition
+            raise_error: If True, raises McpError on validation failure.
+
+        Returns:
+            None if valid, ErrorData if invalid and raise_error=False
+        """
+        mode = tool.execution.task if tool.execution else None
+        return self.validate_task_mode(mode, raise_error=raise_error)
+
+    def can_use_tool(self, tool_task_mode: TaskExecutionMode | None) -> bool:
+        """
+        Check if this client can use a tool with the given task mode.
+
+        Useful for filtering tool lists or providing warnings.
+        Returns False if tool requires "always" but client doesn't support tasks.
+
+        Args:
+            tool_task_mode: The tool's execution.task value
+
+        Returns:
+            True if the client can use this tool, False otherwise
+        """
+        mode = tool_task_mode or "never"
+        if mode == "always" and not self.client_supports_tasks:
+            return False
+        return True
+
 
 @dataclass
 class RequestContext(Generic[SessionT, LifespanContextT, RequestT]):

src/mcp/types.py

@@ -39,7 +39,7 @@
 Role = Literal["user", "assistant"]
 RequestId = Annotated[int, Field(strict=True)] | str
 AnyFunction: TypeAlias = Callable[..., Any]
-TaskHint = Literal["never", "optional", "always"]
+TaskExecutionMode = Literal["never", "optional", "always"]
 
 
 class TaskMetadata(BaseModel):
@@ -1126,7 +1126,7 @@ class ToolExecution(BaseModel):
 
     model_config = ConfigDict(extra="allow")
 
-    task: Literal["never", "optional", "always"] | None = None
+    task: TaskExecutionMode | None = None
     """
     Indicates whether this tool supports task-augmented execution.
     This allows clients to handle long-running operations through polling

tests/shared/test_context.py

@@ -0,0 +1,177 @@
+"""Tests for the RequestContext and Experimental classes."""
+
+import pytest
+
+from mcp.shared.context import Experimental
+from mcp.shared.exceptions import McpError
+from mcp.types import (
+    METHOD_NOT_FOUND,
+    ClientCapabilities,
+    ClientTasksCapability,
+    TaskMetadata,
+    Tool,
+    ToolExecution,
+)
+
+# --- Experimental.is_task ---
+
+
+def test_is_task_true_when_metadata_present() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    assert exp.is_task is True
+
+
+def test_is_task_false_when_no_metadata() -> None:
+    exp = Experimental(task_metadata=None)
+    assert exp.is_task is False
+
+
+# --- Experimental.client_supports_tasks ---
+
+
+def test_client_supports_tasks_true() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities(tasks=ClientTasksCapability()))
+    assert exp.client_supports_tasks is True
+
+
+def test_client_supports_tasks_false_no_tasks() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities())
+    assert exp.client_supports_tasks is False
+
+
+def test_client_supports_tasks_false_no_capabilities() -> None:
+    exp = Experimental(_client_capabilities=None)
+    assert exp.client_supports_tasks is False
+
+
+# --- Experimental.validate_task_mode ---
+
+
+def test_validate_task_mode_always_with_task_is_valid() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    error = exp.validate_task_mode("always", raise_error=False)
+    assert error is None
+
+
+def test_validate_task_mode_always_without_task_returns_error() -> None:
+    exp = Experimental(task_metadata=None)
+    error = exp.validate_task_mode("always", raise_error=False)
+    assert error is not None
+    assert error.code == METHOD_NOT_FOUND
+    assert "requires task-augmented" in error.message
+
+
+def test_validate_task_mode_always_without_task_raises_by_default() -> None:
+    exp = Experimental(task_metadata=None)
+    with pytest.raises(McpError) as exc_info:
+        exp.validate_task_mode("always")
+    assert exc_info.value.error.code == METHOD_NOT_FOUND
+
+
+def test_validate_task_mode_never_without_task_is_valid() -> None:
+    exp = Experimental(task_metadata=None)
+    error = exp.validate_task_mode("never", raise_error=False)
+    assert error is None
+
+
+def test_validate_task_mode_never_with_task_returns_error() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    error = exp.validate_task_mode("never", raise_error=False)
+    assert error is not None
+    assert error.code == METHOD_NOT_FOUND
+    assert "does not support task-augmented" in error.message
+
+
+def test_validate_task_mode_never_with_task_raises_by_default() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    with pytest.raises(McpError) as exc_info:
+        exp.validate_task_mode("never")
+    assert exc_info.value.error.code == METHOD_NOT_FOUND
+
+
+def test_validate_task_mode_none_treated_as_never() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    error = exp.validate_task_mode(None, raise_error=False)
+    assert error is not None
+    assert "does not support task-augmented" in error.message
+
+
+def test_validate_task_mode_optional_with_task_is_valid() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    error = exp.validate_task_mode("optional", raise_error=False)
+    assert error is None
+
+
+def test_validate_task_mode_optional_without_task_is_valid() -> None:
+    exp = Experimental(task_metadata=None)
+    error = exp.validate_task_mode("optional", raise_error=False)
+    assert error is None
+
+
+# --- Experimental.validate_for_tool ---
+
+
+def test_validate_for_tool_with_execution_always() -> None:
+    exp = Experimental(task_metadata=None)
+    tool = Tool(
+        name="test",
+        description="test",
+        inputSchema={"type": "object"},
+        execution=ToolExecution(task="always"),
+    )
+    error = exp.validate_for_tool(tool, raise_error=False)
+    assert error is not None
+    assert "requires task-augmented" in error.message
+
+
+def test_validate_for_tool_without_execution() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    tool = Tool(
+        name="test",
+        description="test",
+        inputSchema={"type": "object"},
+        execution=None,
+    )
+    error = exp.validate_for_tool(tool, raise_error=False)
+    assert error is not None
+    assert "does not support task-augmented" in error.message
+
+
+def test_validate_for_tool_optional_with_task() -> None:
+    exp = Experimental(task_metadata=TaskMetadata(ttl=60000))
+    tool = Tool(
+        name="test",
+        description="test",
+        inputSchema={"type": "object"},
+        execution=ToolExecution(task="optional"),
+    )
+    error = exp.validate_for_tool(tool, raise_error=False)
+    assert error is None
+
+
+# --- Experimental.can_use_tool ---
+
+
+def test_can_use_tool_always_with_task_support() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities(tasks=ClientTasksCapability()))
+    assert exp.can_use_tool("always") is True
+
+
+def test_can_use_tool_always_without_task_support() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities())
+    assert exp.can_use_tool("always") is False
+
+
+def test_can_use_tool_optional_without_task_support() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities())
+    assert exp.can_use_tool("optional") is True
+
+
+def test_can_use_tool_never_without_task_support() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities())
+    assert exp.can_use_tool("never") is True
+
+
+def test_can_use_tool_none_without_task_support() -> None:
+    exp = Experimental(_client_capabilities=ClientCapabilities())
+    assert exp.can_use_tool(None) is True