99from collections .abc import Awaitable , Callable
1010from typing import TYPE_CHECKING
1111
12+ from mcp .server .experimental .task_scope import task_in_session_scope , task_listable_in_session_scope
1213from mcp .server .experimental .task_support import TaskSupport
1314from mcp .server .lowlevel .func_inspection import create_call_wrapper
1415from mcp .shared .exceptions import McpError
3132 ServerResult ,
3233 ServerTasksCapability ,
3334 ServerTasksRequestsCapability ,
35+ Task ,
3436 TasksCallCapability ,
3537 TasksCancelCapability ,
3638 TasksListCapability ,
@@ -125,15 +127,46 @@ def enable_tasks(
125127
126128 return self ._task_support
127129
130+ def _requestor_session_scope (self ) -> str | None :
131+ """Return the task session scope of the session making the current request."""
132+ return self ._server .request_context .session .experimental .task_session_scope
133+
134+ def _require_task_in_requestor_scope (self , task_id : str ) -> None :
135+ """Reject task IDs that belong to a different session.
136+
137+ Task IDs generated by `run_task()` embed the creating session's
138+ scope. The default handlers treat a task created by another session
139+ exactly like a task that does not exist, so a requestor cannot tell
140+ whether such a task exists. Task IDs without an embedded scope are
141+ accepted from any session.
142+
143+ Raises:
144+ McpError: With INVALID_PARAMS if the task belongs to another session.
145+ """
146+ if not task_in_session_scope (task_id , self ._requestor_session_scope ()):
147+ raise McpError (
148+ ErrorData (
149+ code = INVALID_PARAMS ,
150+ message = f"Task not found: { task_id } " ,
151+ )
152+ )
153+
128154 def _register_default_task_handlers (self ) -> None :
129- """Register default handlers for task operations."""
155+ """Register default handlers for task operations.
156+
157+ Each default handler only operates on tasks created by the requesting
158+ session (see `_require_task_in_requestor_scope`), and tasks/list only
159+ returns the requesting session's own tasks (see
160+ `task_listable_in_session_scope`).
161+ """
130162 assert self ._task_support is not None
131163 support = self ._task_support
132164
133165 # Register get_task handler if not already registered
134166 if GetTaskRequest not in self ._request_handlers :
135167
136168 async def _default_get_task (req : GetTaskRequest ) -> ServerResult :
169+ self ._require_task_in_requestor_scope (req .params .taskId )
137170 task = await support .store .get_task (req .params .taskId )
138171 if task is None :
139172 raise McpError (
@@ -160,6 +193,7 @@ async def _default_get_task(req: GetTaskRequest) -> ServerResult:
160193 if GetTaskPayloadRequest not in self ._request_handlers :
161194
162195 async def _default_get_task_result (req : GetTaskPayloadRequest ) -> ServerResult :
196+ self ._require_task_in_requestor_scope (req .params .taskId )
163197 ctx = self ._server .request_context
164198 result = await support .handler .handle (req , ctx .session , ctx .request_id )
165199 return ServerResult (result )
@@ -170,16 +204,34 @@ async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult:
170204 if ListTasksRequest not in self ._request_handlers :
171205
172206 async def _default_list_tasks (req : ListTasksRequest ) -> ServerResult :
173- cursor = req .params .cursor if req .params else None
174- tasks , next_cursor = await support .store .list_tasks (cursor )
175- return ServerResult (ListTasksResult (tasks = tasks , nextCursor = next_cursor ))
207+ requestor_scope = self ._requestor_session_scope ()
208+ if requestor_scope is None :
209+ # The server cannot tell this requestor apart from any
210+ # other, so there are no tasks it can be shown.
211+ return ServerResult (ListTasksResult (tasks = []))
212+ # Return every task that belongs to the requesting session in
213+ # a single page. The store's pagination cursor is never sent
214+ # to the requestor: it is derived from the unfiltered listing,
215+ # so it could identify a task belonging to a different
216+ # session. For the same reason the request's cursor is not
217+ # forwarded to the store.
218+ own_tasks : list [Task ] = []
219+ cursor : str | None = None
220+ while True :
221+ page , cursor = await support .store .list_tasks (cursor )
222+ own_tasks .extend (
223+ task for task in page if task_listable_in_session_scope (task .taskId , requestor_scope )
224+ )
225+ if cursor is None :
226+ return ServerResult (ListTasksResult (tasks = own_tasks ))
176227
177228 self ._request_handlers [ListTasksRequest ] = _default_list_tasks
178229
179230 # Register cancel_task handler if not already registered
180231 if CancelTaskRequest not in self ._request_handlers :
181232
182233 async def _default_cancel_task (req : CancelTaskRequest ) -> ServerResult :
234+ self ._require_task_in_requestor_scope (req .params .taskId )
183235 result = await cancel_task (support .store , req .params .taskId )
184236 return ServerResult (result )
185237
0 commit comments