Add cancel_task helper and terminal status transition validation

Add spec-compliant task cancellation and status transition handling:

- Add cancel_task() helper that validates task state before cancellation,
  returning -32602 (Invalid params) for nonexistent or terminal tasks
- Add terminal status transition validation in InMemoryTaskStore.update_task()
  to prevent transitions from completed/failed/cancelled states
- Export cancel_task from the tasks module for easy access
- Add comprehensive tests for both features

Per spec: "Receivers MUST reject cancellation of terminal status tasks
with -32602 (Invalid params)" and "Terminal states MUST NOT transition
to any other status"

Author: maxisbey

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

@@ -7,7 +7,7 @@
 - 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
+- Helper functions: run_task, is_terminal, create_task_state, generate_task_id, cancel_task
 
 Architecture:
 - TaskStore is pure storage - it doesn't know about execution
@@ -20,6 +20,7 @@
 
 from mcp.shared.experimental.tasks.context import TaskContext
 from mcp.shared.experimental.tasks.helpers import (
+    cancel_task,
     create_task_state,
     generate_task_id,
     is_terminal,
@@ -53,4 +54,5 @@
     "is_terminal",
     "create_task_state",
     "generate_task_id",
+    "cancel_task",
 ]

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

@@ -10,9 +10,19 @@
 
 from anyio.abc import TaskGroup
 
+from mcp.shared.exceptions import McpError
 from mcp.shared.experimental.tasks.context import TaskContext
 from mcp.shared.experimental.tasks.store import TaskStore
-from mcp.types import CreateTaskResult, Result, Task, TaskMetadata, TaskStatus
+from mcp.types import (
+    INVALID_PARAMS,
+    CancelTaskResult,
+    CreateTaskResult,
+    ErrorData,
+    Result,
+    Task,
+    TaskMetadata,
+    TaskStatus,
+)
 
 if TYPE_CHECKING:
     from mcp.server.session import ServerSession
@@ -33,6 +43,58 @@ def is_terminal(status: TaskStatus) -> bool:
     return status in ("completed", "failed", "cancelled")
 
 
+async def cancel_task(
+    store: TaskStore,
+    task_id: str,
+) -> CancelTaskResult:
+    """
+    Cancel a task with spec-compliant validation.
+
+    Per spec: "Receivers MUST reject cancellation of terminal status tasks
+    with -32602 (Invalid params)"
+
+    This helper validates that the task exists and is not in a terminal state
+    before setting it to "cancelled".
+
+    Args:
+        store: The task store
+        task_id: The task identifier to cancel
+
+    Returns:
+        CancelTaskResult with the cancelled task state
+
+    Raises:
+        McpError: With INVALID_PARAMS (-32602) if:
+            - Task does not exist
+            - Task is already in a terminal state (completed, failed, cancelled)
+
+    Example:
+        @server.experimental.cancel_task()
+        async def handle_cancel(request: CancelTaskRequest) -> CancelTaskResult:
+            return await cancel_task(store, request.params.taskId)
+    """
+    task = await store.get_task(task_id)
+    if task is None:
+        raise McpError(
+            ErrorData(
+                code=INVALID_PARAMS,
+                message=f"Task not found: {task_id}",
+            )
+        )
+
+    if is_terminal(task.status):
+        raise McpError(
+            ErrorData(
+                code=INVALID_PARAMS,
+                message=f"Cannot cancel task in terminal state '{task.status}'",
+            )
+        )
+
+    # Update task to cancelled status
+    cancelled_task = await store.update_task(task_id, status="cancelled")
+    return CancelTaskResult(**cancelled_task.model_dump())
+
+
 def generate_task_id() -> str:
     """Generate a unique task ID."""
     return str(uuid4())

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

@@ -114,6 +114,10 @@ async def update_task(
         if stored is None:
             raise ValueError(f"Task with ID {task_id} not found")
 
+        # Per spec: Terminal states MUST NOT transition to any other status
+        if status is not None and status != stored.task.status and is_terminal(stored.task.status):
+            raise ValueError(f"Cannot transition from terminal status '{stored.task.status}'")
+
         status_changed = False
         if status is not None and stored.task.status != status:
             stored.task.status = status

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

@@ -69,6 +69,9 @@ async def update_task(
 
         Raises:
             ValueError: If task not found
+            ValueError: If attempting to transition from a terminal status
+                (completed, failed, cancelled). Per spec, terminal states
+                MUST NOT transition to any other status.
         """
 
     @abstractmethod

tests/experimental/tasks/server/test_store.py

@@ -4,8 +4,9 @@
 
 import pytest
 
-from mcp.shared.experimental.tasks import InMemoryTaskStore
-from mcp.types import CallToolResult, TaskMetadata, TextContent
+from mcp.shared.exceptions import McpError
+from mcp.shared.experimental.tasks import InMemoryTaskStore, cancel_task
+from mcp.types import INVALID_PARAMS, CallToolResult, TaskMetadata, TextContent
 
 
 @pytest.mark.anyio
@@ -339,3 +340,160 @@ async def test_terminal_task_ttl_reset() -> None:
     assert new_expiry >= initial_expiry
 
     store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_terminal_status_transition_rejected() -> None:
+    """Test that transitions from terminal states are rejected.
+
+    Per spec: Terminal states (completed, failed, cancelled) MUST NOT
+    transition to any other status.
+    """
+    store = InMemoryTaskStore()
+
+    # Test each terminal status
+    for terminal_status in ("completed", "failed", "cancelled"):
+        task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+
+        # Move to terminal state
+        await store.update_task(task.taskId, status=terminal_status)
+
+        # Attempting to transition to any other status should raise
+        with pytest.raises(ValueError, match="Cannot transition from terminal status"):
+            await store.update_task(task.taskId, status="working")
+
+        # Also test transitioning to another terminal state
+        other_terminal = "failed" if terminal_status != "failed" else "completed"
+        with pytest.raises(ValueError, match="Cannot transition from terminal status"):
+            await store.update_task(task.taskId, status=other_terminal)
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_terminal_status_allows_same_status() -> None:
+    """Test that setting the same terminal status doesn't raise.
+
+    This is not a transition, so it should be allowed (no-op).
+    """
+    store = InMemoryTaskStore()
+
+    task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+    await store.update_task(task.taskId, status="completed")
+
+    # Setting the same status should not raise
+    updated = await store.update_task(task.taskId, status="completed")
+    assert updated.status == "completed"
+
+    # Updating just the message should also work
+    updated = await store.update_task(task.taskId, status_message="Updated message")
+    assert updated.statusMessage == "Updated message"
+
+    store.cleanup()
+
+
+# =============================================================================
+# cancel_task helper function tests
+# =============================================================================
+
+
+@pytest.mark.anyio
+async def test_cancel_task_succeeds_for_working_task() -> None:
+    """Test cancel_task helper succeeds for a working task."""
+    store = InMemoryTaskStore()
+
+    task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+    assert task.status == "working"
+
+    result = await cancel_task(store, task.taskId)
+
+    assert result.taskId == task.taskId
+    assert result.status == "cancelled"
+
+    # Verify store is updated
+    retrieved = await store.get_task(task.taskId)
+    assert retrieved is not None
+    assert retrieved.status == "cancelled"
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_cancel_task_rejects_nonexistent_task() -> None:
+    """Test cancel_task raises McpError with INVALID_PARAMS for nonexistent task."""
+    store = InMemoryTaskStore()
+
+    with pytest.raises(McpError) as exc_info:
+        await cancel_task(store, "nonexistent-task-id")
+
+    assert exc_info.value.error.code == INVALID_PARAMS
+    assert "not found" in exc_info.value.error.message
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_cancel_task_rejects_completed_task() -> None:
+    """Test cancel_task raises McpError with INVALID_PARAMS for completed task."""
+    store = InMemoryTaskStore()
+
+    task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+    await store.update_task(task.taskId, status="completed")
+
+    with pytest.raises(McpError) as exc_info:
+        await cancel_task(store, task.taskId)
+
+    assert exc_info.value.error.code == INVALID_PARAMS
+    assert "terminal state 'completed'" in exc_info.value.error.message
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_cancel_task_rejects_failed_task() -> None:
+    """Test cancel_task raises McpError with INVALID_PARAMS for failed task."""
+    store = InMemoryTaskStore()
+
+    task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+    await store.update_task(task.taskId, status="failed")
+
+    with pytest.raises(McpError) as exc_info:
+        await cancel_task(store, task.taskId)
+
+    assert exc_info.value.error.code == INVALID_PARAMS
+    assert "terminal state 'failed'" in exc_info.value.error.message
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_cancel_task_rejects_already_cancelled_task() -> None:
+    """Test cancel_task raises McpError with INVALID_PARAMS for already cancelled task."""
+    store = InMemoryTaskStore()
+
+    task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+    await store.update_task(task.taskId, status="cancelled")
+
+    with pytest.raises(McpError) as exc_info:
+        await cancel_task(store, task.taskId)
+
+    assert exc_info.value.error.code == INVALID_PARAMS
+    assert "terminal state 'cancelled'" in exc_info.value.error.message
+
+    store.cleanup()
+
+
+@pytest.mark.anyio
+async def test_cancel_task_succeeds_for_input_required_task() -> None:
+    """Test cancel_task helper succeeds for a task in input_required status."""
+    store = InMemoryTaskStore()
+
+    task = await store.create_task(metadata=TaskMetadata(ttl=60000))
+    await store.update_task(task.taskId, status="input_required")
+
+    result = await cancel_task(store, task.taskId)
+
+    assert result.taskId == task.taskId
+    assert result.status == "cancelled"
+
+    store.cleanup()