Skip to content

Commit 6213787

Browse files
authored
[v1.x] Scope experimental tasks to the session that created them (#2720)
1 parent ce267b6 commit 6213787

13 files changed

Lines changed: 731 additions & 14 deletions

File tree

docs/experimental/tasks-server.md

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,29 @@ That's it. `enable_tasks()` automatically:
5353
- Registers handlers for `tasks/get`, `tasks/result`, `tasks/list`, `tasks/cancel`
5454
- Updates server capabilities
5555

56+
## Task Visibility
57+
58+
Task IDs generated by `run_task()` embed an opaque marker identifying the session that
59+
created the task, and the default handlers use it to restrict each session to its own
60+
tasks: `tasks/get`, `tasks/result`, and `tasks/cancel` respond with "task not found" for
61+
another session's task, and `tasks/list` returns only the requesting session's tasks. A
62+
client that reconnects gets a new session and can no longer reach tasks it created on the
63+
previous one.
64+
65+
A task ID has no session marker when it was passed to `run_task()` explicitly, when the
66+
task was created directly through the `TaskStore`, or when the server runs in stateless
67+
mode (each request gets a fresh session, so tasks must remain reachable across requests).
68+
Such tasks are accessible to any requestor that presents the exact task ID, and are never
69+
included in `tasks/list` responses because the server cannot tell which session they
70+
belong to. Treat these task IDs as capabilities: generate them with enough entropy that
71+
they cannot be guessed, share them only with the intended recipient, and prefer short
72+
TTLs. Passing an explicit `task_id` to `run_task()` is deprecated for this reason.
73+
74+
To scope tasks to something other than the session — for example a user identity from your
75+
authorization layer — register your own handlers with `@server.experimental.get_task()`,
76+
`@server.experimental.get_task_result()`, `@server.experimental.list_tasks()`, and
77+
`@server.experimental.cancel_task()` instead of relying on the defaults.
78+
5679
## Tool Declaration
5780

5881
Tools declare task support via the `execution.taskSupport` field:

src/mcp/server/experimental/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,5 @@
88
- mcp.server.experimental.task_support.TaskSupport
99
- mcp.server.experimental.task_result_handler.TaskResultHandler
1010
- mcp.server.experimental.request_context.Experimental
11+
- mcp.server.experimental.task_scope (session scoping of task IDs)
1112
"""

src/mcp/server/experimental/request_context.py

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,15 @@
77
WARNING: These APIs are experimental and may change without notice.
88
"""
99

10+
import warnings
1011
from collections.abc import Awaitable, Callable
1112
from dataclasses import dataclass, field
12-
from typing import Any
13+
from typing import Any, overload
14+
15+
from typing_extensions import deprecated
1316

1417
from mcp.server.experimental.task_context import ServerTaskContext
18+
from mcp.server.experimental.task_scope import scoped_task_id
1519
from mcp.server.experimental.task_support import TaskSupport
1620
from mcp.server.session import ServerSession
1721
from mcp.shared.exceptions import McpError
@@ -29,6 +33,14 @@
2933
Tool,
3034
)
3135

36+
EXPLICIT_TASK_ID_DEPRECATION = (
37+
"Passing an explicit task_id to run_task is deprecated. A task created with an "
38+
"explicit ID is not associated with the session that created it: any requestor "
39+
"that presents the ID can read its status and result or cancel it, and it never "
40+
"appears in tasks/list. Omit task_id to let the SDK generate an ID associated "
41+
"with the creating session."
42+
)
43+
3244

3345
@dataclass
3446
class Experimental:
@@ -143,6 +155,25 @@ def can_use_tool(self, tool_task_mode: TaskExecutionMode | None) -> bool:
143155
return False
144156
return True
145157

158+
@overload
159+
async def run_task(
160+
self,
161+
work: Callable[[ServerTaskContext], Awaitable[Result]],
162+
*,
163+
task_id: None = None,
164+
model_immediate_response: str | None = None,
165+
) -> CreateTaskResult: ...
166+
167+
@overload
168+
@deprecated(EXPLICIT_TASK_ID_DEPRECATION)
169+
async def run_task(
170+
self,
171+
work: Callable[[ServerTaskContext], Awaitable[Result]],
172+
*,
173+
task_id: str,
174+
model_immediate_response: str | None = None,
175+
) -> CreateTaskResult: ...
176+
146177
async def run_task(
147178
self,
148179
work: Callable[[ServerTaskContext], Awaitable[Result]],
@@ -167,9 +198,17 @@ async def run_task(
167198
When work() returns a Result, the task is auto-completed with that result.
168199
If work() raises an exception, the task is auto-failed.
169200
201+
Generated task IDs embed the session's task scope so that the default
202+
task handlers only serve the task to the session that created it. An
203+
explicitly provided `task_id` is used verbatim and is not associated
204+
with the session, so any session can access it through the default
205+
handlers; passing one is deprecated for that reason.
206+
170207
Args:
171208
work: Async function that does the actual work
172-
task_id: Optional task ID (generated if not provided)
209+
task_id: Deprecated. Optional task ID, used verbatim and not
210+
associated with the creating session. Omit it to let the SDK
211+
generate one.
173212
model_immediate_response: Optional string to include in _meta as
174213
io.modelcontextprotocol/model-immediate-response
175214
@@ -196,6 +235,8 @@ async def work(task: ServerTaskContext) -> CallToolResult:
196235
197236
WARNING: This API is experimental and may change without notice.
198237
"""
238+
if task_id is not None:
239+
warnings.warn(EXPLICIT_TASK_ID_DEPRECATION, DeprecationWarning, stacklevel=2)
199240
if self._task_support is None:
200241
raise RuntimeError("Task support not enabled. Call server.experimental.enable_tasks() first.")
201242
if self._session is None:
@@ -210,6 +251,11 @@ async def work(task: ServerTaskContext) -> CallToolResult:
210251
# Access task_group via TaskSupport - raises if not in run() context
211252
task_group = support.task_group
212253

254+
if task_id is None:
255+
session_scope = self._session.experimental.task_session_scope
256+
if session_scope is not None:
257+
task_id = scoped_task_id(session_scope)
258+
213259
task = await support.store.create_task(self.task_metadata, task_id)
214260

215261
task_ctx = ServerTaskContext(

src/mcp/server/experimental/session_features.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,12 @@ class ExperimentalServerSessionFeatures:
4040

4141
def __init__(self, session: "ServerSession") -> None:
4242
self._session = session
43+
# Opaque marker identifying this session for task scoping. Assigned by
44+
# TaskSupport.configure_session(). Task IDs generated by run_task()
45+
# embed it so the default task handlers can restrict task access to
46+
# the session that created the task. None means tasks created on this
47+
# session are not associated with it (e.g. stateless servers).
48+
self.task_session_scope: str | None = None
4349

4450
async def get_task(self, task_id: str) -> types.GetTaskResult:
4551
"""

src/mcp/server/experimental/task_result_handler.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,11 @@ class TaskResultHandler:
4646
4. Blocks until task reaches terminal state
4747
5. Returns the final result
4848
49+
Prefer `server.experimental.enable_tasks()`, whose default tasks/result
50+
handler wraps `handle()` and only serves tasks created by the requesting
51+
session. A custom handler that calls `handle()` directly is responsible
52+
for deciding which requestors may access which tasks.
53+
4954
Usage:
5055
# Create handler with store and queue
5156
handler = TaskResultHandler(task_store, message_queue)
@@ -55,9 +60,6 @@ class TaskResultHandler:
5560
async def handle_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
5661
ctx = server.request_context
5762
return await handler.handle(req, ctx.session, ctx.request_id)
58-
59-
# Or use the convenience method
60-
handler.register(server)
6163
"""
6264

6365
def __init__(
Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
"""
2+
Session scoping for experimental task identifiers.
3+
4+
Task IDs generated by `run_task()` embed an opaque, per-session marker (the
5+
"session scope") so that the default task handlers can tell which session
6+
created a task. The default handlers for tasks/get, tasks/result, tasks/list,
7+
and tasks/cancel only operate on tasks created by the requesting session.
8+
9+
Task IDs without a session scope (explicitly provided IDs, IDs created
10+
directly through a TaskStore, or IDs created in stateless mode) have no known
11+
creator. They can be used with tasks/get, tasks/result, and tasks/cancel from
12+
any session - possession of the ID is what grants access - but they are never
13+
included in tasks/list responses.
14+
15+
WARNING: These APIs are experimental and may change without notice.
16+
"""
17+
18+
import re
19+
from uuid import uuid4
20+
21+
__all__ = [
22+
"new_session_scope",
23+
"scoped_task_id",
24+
"session_scope_of",
25+
"task_in_session_scope",
26+
"task_listable_in_session_scope",
27+
]
28+
29+
# A scoped task ID has the form "<32 hex chars>:<uuid4>". Both halves must
30+
# match exactly so that explicitly chosen task IDs are never mistaken for
31+
# scoped ones. \Z rather than $ so a trailing newline cannot match.
32+
_SCOPED_TASK_ID = re.compile(
33+
r"\A(?P<scope>[0-9a-f]{32}):"
34+
r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\Z"
35+
)
36+
37+
38+
def new_session_scope() -> str:
39+
"""Create a new opaque session scope token."""
40+
return uuid4().hex
41+
42+
43+
def scoped_task_id(session_scope: str) -> str:
44+
"""Generate a task ID associated with the given session scope."""
45+
return f"{session_scope}:{uuid4()}"
46+
47+
48+
def session_scope_of(task_id: str) -> str | None:
49+
"""Return the session scope embedded in a task ID, or None if it has none."""
50+
match = _SCOPED_TASK_ID.match(task_id)
51+
return match.group("scope") if match else None
52+
53+
54+
def task_in_session_scope(task_id: str, session_scope: str | None) -> bool:
55+
"""Whether a task may be used by a requestor with the given session scope.
56+
57+
Used by tasks/get, tasks/result, and tasks/cancel. A task whose ID carries
58+
no session scope has no known creator, so possession of the ID is what
59+
grants access to it: it can be used from any session.
60+
"""
61+
embedded = session_scope_of(task_id)
62+
return embedded is None or embedded == session_scope
63+
64+
65+
def task_listable_in_session_scope(task_id: str, session_scope: str | None) -> bool:
66+
"""Whether a task may be included in a tasks/list response for the given session scope.
67+
68+
Used by tasks/list. Listing is stricter than access by ID: a task is only
69+
listed to the session that created it. Tasks with no session scope are
70+
never listed because they have no known creator, and requestors with no
71+
session scope are never shown any tasks because the server cannot tell
72+
them apart.
73+
"""
74+
embedded = session_scope_of(task_id)
75+
return embedded is not None and embedded == session_scope

src/mcp/server/experimental/task_support.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from anyio.abc import TaskGroup
1414

1515
from mcp.server.experimental.task_result_handler import TaskResultHandler
16+
from mcp.server.experimental.task_scope import new_session_scope
1617
from mcp.server.session import ServerSession
1718
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
1819
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
@@ -83,20 +84,30 @@ async def run(self) -> AsyncIterator[None]:
8384
finally:
8485
self._task_group = None
8586

86-
def configure_session(self, session: ServerSession) -> None:
87+
def configure_session(self, session: ServerSession, *, stateless: bool = False) -> None:
8788
"""
8889
Configure a session for task support.
8990
9091
This registers the result handler as a response router so that
9192
responses to queued requests (elicitation, sampling) are routed
9293
back to the waiting resolvers.
9394
95+
It also assigns the session a task session scope. Task IDs generated
96+
by `run_task()` embed this scope, and the default task handlers only
97+
operate on tasks created by the requesting session. Stateless sessions
98+
are not assigned a scope: each request runs on a fresh session, so a
99+
task created by one request could never be retrieved by a later one if
100+
tasks were bound to the session that created them.
101+
94102
Called automatically by Server.run() for each new session.
95103
96104
Args:
97105
session: The session to configure
106+
stateless: Whether the session belongs to a stateless server run
98107
"""
99108
session.add_response_router(self.handler)
109+
if not stateless and session.experimental.task_session_scope is None:
110+
session.experimental.task_session_scope = new_session_scope()
100111

101112
@classmethod
102113
def in_memory(cls) -> "TaskSupport":

src/mcp/server/lowlevel/experimental.py

Lines changed: 56 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from collections.abc import Awaitable, Callable
1010
from typing import TYPE_CHECKING
1111

12+
from mcp.server.experimental.task_scope import task_in_session_scope, task_listable_in_session_scope
1213
from mcp.server.experimental.task_support import TaskSupport
1314
from mcp.server.lowlevel.func_inspection import create_call_wrapper
1415
from mcp.shared.exceptions import McpError
@@ -31,6 +32,7 @@
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

Comments
 (0)