Fix the session binding logic for tasks.

Author: localden

src/mcp/server/experimental/request_context.py

@@ -187,7 +187,8 @@ async def work(task: ServerTaskContext) -> CallToolResult:
         # Access task_group via TaskSupport - raises if not in run() context
         task_group = support.task_group
 
-        task = await support.store.create_task(self.task_metadata, task_id)
+        session_id = str(id(self._session))
+        task = await support.store.create_task(self.task_metadata, task_id, session_id=session_id)
 
         task_ctx = ServerTaskContext(
             task=task,

src/mcp/server/experimental/task_result_handler.py

@@ -80,6 +80,7 @@ async def handle(
         request: GetTaskPayloadRequest,
         session: ServerSession,
         request_id: RequestId,
+        session_id: str | None = None,
     ) -> GetTaskPayloadResult:
         """Handle a tasks/result request.
 
@@ -94,22 +95,23 @@ async def handle(
             request: The GetTaskPayloadRequest
             session: The server session for sending messages
             request_id: The request ID for relatedRequestId routing
+            session_id: Optional session identifier for access control.
 
         Returns:
             GetTaskPayloadResult with the task's final payload
         """
         task_id = request.params.task_id
 
         while True:
-            task = await self._store.get_task(task_id)
+            task = await self._store.get_task(task_id, session_id=session_id)
             if task is None:
                 raise MCPError(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)
+                result = await self._store.get_result(task_id, session_id=session_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

src/mcp/server/lowlevel/experimental.py

@@ -153,7 +153,8 @@ def enable_tasks(
             async def _default_get_task(
                 ctx: ServerRequestContext[LifespanResultT], params: GetTaskRequestParams
             ) -> GetTaskResult:
-                task = await task_support.store.get_task(params.task_id)
+                session_id = str(id(ctx.session))
+                task = await task_support.store.get_task(params.task_id, session_id=session_id)
                 if task is None:
                     raise MCPError(code=INVALID_PARAMS, message=f"Task not found: {params.task_id}")
                 return GetTaskResult(
@@ -174,8 +175,9 @@ async def _default_get_task_result(
                 ctx: ServerRequestContext[LifespanResultT], params: GetTaskPayloadRequestParams
             ) -> GetTaskPayloadResult:
                 assert ctx.request_id is not None
+                session_id = str(id(ctx.session))
                 req = GetTaskPayloadRequest(params=params)
-                result = await task_support.handler.handle(req, ctx.session, ctx.request_id)
+                result = await task_support.handler.handle(req, ctx.session, ctx.request_id, session_id=session_id)
                 return result
 
             self._add_request_handler("tasks/result", _default_get_task_result)
@@ -186,7 +188,8 @@ async def _default_list_tasks(
                 ctx: ServerRequestContext[LifespanResultT], params: PaginatedRequestParams | None
             ) -> ListTasksResult:
                 cursor = params.cursor if params else None
-                tasks, next_cursor = await task_support.store.list_tasks(cursor)
+                session_id = str(id(ctx.session))
+                tasks, next_cursor = await task_support.store.list_tasks(cursor, session_id=session_id)
                 return ListTasksResult(tasks=tasks, next_cursor=next_cursor)
 
             self._add_request_handler("tasks/list", _default_list_tasks)
@@ -196,7 +199,8 @@ async def _default_list_tasks(
             async def _default_cancel_task(
                 ctx: ServerRequestContext[LifespanResultT], params: CancelTaskRequestParams
             ) -> CancelTaskResult:
-                result = await cancel_task(task_support.store, params.task_id)
+                session_id = str(id(ctx.session))
+                result = await cancel_task(task_support.store, params.task_id, session_id=session_id)
                 return result
 
             self._add_request_handler("tasks/cancel", _default_cancel_task)

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

@@ -50,6 +50,7 @@ def is_terminal(status: TaskStatus) -> bool:
 async def cancel_task(
     store: TaskStore,
     task_id: str,
+    session_id: str | None = None,
 ) -> CancelTaskResult:
     """Cancel a task with spec-compliant validation.
 
@@ -62,28 +63,29 @@ async def cancel_task(
     Args:
         store: The task store
         task_id: The task identifier to cancel
+        session_id: Optional session identifier for access control.
 
     Returns:
         CancelTaskResult with the cancelled task state
 
     Raises:
         MCPError: With INVALID_PARAMS (-32602) if:
-            - Task does not exist
+            - Task does not exist or is not accessible by this session
             - Task is already in a terminal state (completed, failed, cancelled)
 
     Example:
         async def handle_cancel(ctx, params: CancelTaskRequestParams) -> CancelTaskResult:
             return await cancel_task(store, params.task_id)
     """
-    task = await store.get_task(task_id)
+    task = await store.get_task(task_id, session_id=session_id)
     if task is None:
         raise MCPError(code=INVALID_PARAMS, message=f"Task not found: {task_id}")
 
     if is_terminal(task.status):
         raise MCPError(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=TASK_STATUS_CANCELLED)
+    cancelled_task = await store.update_task(task_id, status=TASK_STATUS_CANCELLED, session_id=session_id)
     return CancelTaskResult(**cancelled_task.model_dump())
 
 

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

@@ -22,6 +22,7 @@ class StoredTask:
     """Internal storage representation of a task."""
 
     task: Task
+    session_id: str | None = None
     result: Result | None = None
     # Time when this task should be removed (None = never)
     expires_at: datetime | None = field(default=None)
@@ -32,6 +33,7 @@ class InMemoryTaskStore(TaskStore):
 
     Features:
     - Automatic TTL-based cleanup (lazy expiration)
+    - Session isolation (tasks are scoped to their creating session)
     - Thread-safe for single-process async use
     - Pagination support for list_tasks
 
@@ -66,10 +68,25 @@ def _cleanup_expired(self) -> None:
         for task_id in expired_ids:
             del self._tasks[task_id]
 
+    def _get_stored_task(self, task_id: str, session_id: str | None = None) -> StoredTask | None:
+        """Retrieve a stored task, enforcing session ownership when a session_id is provided.
+
+        Returns None if the task does not exist or belongs to a different session.
+        When either the caller's session_id or the stored task's session_id is None,
+        no filtering occurs (backward compatibility).
+        """
+        stored = self._tasks.get(task_id)
+        if stored is None:
+            return None
+        if session_id is not None and stored.session_id is not None and stored.session_id != session_id:
+            return None
+        return stored
+
     async def create_task(
         self,
         metadata: TaskMetadata,
         task_id: str | None = None,
+        session_id: str | None = None,
     ) -> Task:
         """Create a new task with the given metadata."""
         # Cleanup expired tasks on access
@@ -82,19 +99,20 @@ async def create_task(
 
         stored = StoredTask(
             task=task,
+            session_id=session_id,
             expires_at=self._calculate_expiry(metadata.ttl),
         )
         self._tasks[task.task_id] = stored
 
         # Return a copy to prevent external modification
         return Task(**task.model_dump())
 
-    async def get_task(self, task_id: str) -> Task | None:
+    async def get_task(self, task_id: str, session_id: str | None = None) -> Task | None:
         """Get a task by ID."""
         # Cleanup expired tasks on access
         self._cleanup_expired()
 
-        stored = self._tasks.get(task_id)
+        stored = self._get_stored_task(task_id, session_id)
         if stored is None:
             return None
 
@@ -106,9 +124,10 @@ async def update_task(
         task_id: str,
         status: TaskStatus | None = None,
         status_message: str | None = None,
+        session_id: str | None = None,
     ) -> Task:
         """Update a task's status and/or message."""
-        stored = self._tasks.get(task_id)
+        stored = self._get_stored_task(task_id, session_id)
         if stored is None:
             raise ValueError(f"Task with ID {task_id} not found")
 
@@ -137,17 +156,17 @@ async def update_task(
 
         return Task(**stored.task.model_dump())
 
-    async def store_result(self, task_id: str, result: Result) -> None:
+    async def store_result(self, task_id: str, result: Result, session_id: str | None = None) -> None:
         """Store the result for a task."""
-        stored = self._tasks.get(task_id)
+        stored = self._get_stored_task(task_id, session_id)
         if stored is None:
             raise ValueError(f"Task with ID {task_id} not found")
 
         stored.result = result
 
-    async def get_result(self, task_id: str) -> Result | None:
+    async def get_result(self, task_id: str, session_id: str | None = None) -> Result | None:
         """Get the stored result for a task."""
-        stored = self._tasks.get(task_id)
+        stored = self._get_stored_task(task_id, session_id)
         if stored is None:
             return None
 
@@ -156,34 +175,41 @@ async def get_result(self, task_id: str) -> Result | None:
     async def list_tasks(
         self,
         cursor: str | None = None,
+        session_id: str | None = None,
     ) -> tuple[list[Task], str | None]:
         """List tasks with pagination."""
         # Cleanup expired tasks on access
         self._cleanup_expired()
 
-        all_task_ids = list(self._tasks.keys())
+        # Filter tasks by session ownership before pagination
+        filtered_task_ids = [
+            task_id
+            for task_id, stored in self._tasks.items()
+            if session_id is None or stored.session_id is None or stored.session_id == session_id
+        ]
 
         start_index = 0
         if cursor is not None:
             try:
-                cursor_index = all_task_ids.index(cursor)
+                cursor_index = filtered_task_ids.index(cursor)
                 start_index = cursor_index + 1
             except ValueError:
                 raise ValueError(f"Invalid cursor: {cursor}")
 
-        page_task_ids = all_task_ids[start_index : start_index + self._page_size]
+        page_task_ids = filtered_task_ids[start_index : start_index + self._page_size]
         tasks = [Task(**self._tasks[tid].task.model_dump()) for tid in page_task_ids]
 
         # Determine next cursor
         next_cursor = None
-        if start_index + self._page_size < len(all_task_ids) and page_task_ids:
+        if start_index + self._page_size < len(filtered_task_ids) and page_task_ids:
             next_cursor = page_task_ids[-1]
 
         return tasks, next_cursor
 
-    async def delete_task(self, task_id: str) -> bool:
+    async def delete_task(self, task_id: str, session_id: str | None = None) -> bool:
         """Delete a task."""
-        if task_id not in self._tasks:
+        stored = self._get_stored_task(task_id, session_id)
+        if stored is None:
             return False
 
         del self._tasks[task_id]

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

@@ -19,12 +19,15 @@ async def create_task(
         self,
         metadata: TaskMetadata,
         task_id: str | None = None,
+        session_id: str | None = None,
     ) -> Task:
         """Create a new task.
 
         Args:
             metadata: Task metadata (ttl, etc.)
             task_id: Optional task ID. If None, implementation should generate one.
+            session_id: Optional session identifier. When provided, the task is
+                bound to this session for isolation purposes.
 
         Returns:
             The created Task with status="working"
@@ -34,14 +37,15 @@ async def create_task(
         """
 
     @abstractmethod
-    async def get_task(self, task_id: str) -> Task | None:
+    async def get_task(self, task_id: str, session_id: str | None = None) -> Task | None:
         """Get a task by ID.
 
         Args:
             task_id: The task identifier
+            session_id: Optional session identifier for access control.
 
         Returns:
-            The Task, or None if not found
+            The Task, or None if not found or not accessible by this session.
         """
 
     @abstractmethod
@@ -50,70 +54,78 @@ async def update_task(
         task_id: str,
         status: TaskStatus | None = None,
         status_message: str | None = None,
+        session_id: str | None = None,
     ) -> Task:
         """Update a task's status and/or message.
 
         Args:
             task_id: The task identifier
             status: New status (if changing)
             status_message: New status message (if changing)
+            session_id: Optional session identifier for access control.
 
         Returns:
             The updated Task
 
         Raises:
-            ValueError: If task not found
+            ValueError: If task not found or not accessible by this session.
             ValueError: If attempting to transition from a terminal status
                 (completed, failed, cancelled). Per spec, terminal states
                 MUST NOT transition to any other status.
         """
 
     @abstractmethod
-    async def store_result(self, task_id: str, result: Result) -> None:
+    async def store_result(self, task_id: str, result: Result, session_id: str | None = None) -> None:
         """Store the result for a task.
 
         Args:
             task_id: The task identifier
             result: The result to store
+            session_id: Optional session identifier for access control.
 
         Raises:
-            ValueError: If task not found
+            ValueError: If task not found or not accessible by this session.
         """
 
     @abstractmethod
-    async def get_result(self, task_id: str) -> Result | None:
+    async def get_result(self, task_id: str, session_id: str | None = None) -> Result | None:
         """Get the stored result for a task.
 
         Args:
             task_id: The task identifier
+            session_id: Optional session identifier for access control.
 
         Returns:
-            The stored Result, or None if not available
+            The stored Result, or None if not available.
         """
 
     @abstractmethod
     async def list_tasks(
         self,
         cursor: str | None = None,
+        session_id: str | None = None,
     ) -> tuple[list[Task], str | None]:
         """List tasks with pagination.
 
         Args:
             cursor: Optional cursor for pagination
+            session_id: Optional session identifier. When provided, only tasks
+                belonging to this session are returned.
 
         Returns:
             Tuple of (tasks, next_cursor). next_cursor is None if no more pages.
         """
 
     @abstractmethod
-    async def delete_task(self, task_id: str) -> bool:
+    async def delete_task(self, task_id: str, session_id: str | None = None) -> bool:
         """Delete a task.
 
         Args:
             task_id: The task identifier
+            session_id: Optional session identifier for access control.
 
         Returns:
-            True if deleted, False if not found
+            True if deleted, False if not found or not accessible by this session.
         """
 
     @abstractmethod