[v1.x] Scope experimental tasks to the session that created them (#2720)
Some checks failed
Main branch checks / checks (push) Has been cancelled

This commit is contained in:
Max 2026-05-29 18:05:58 +01:00 • committed by GitHub
commit 62137874ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 731 additions and 14 deletions

View file

@ -8,4 +8,5 @@ Import directly from submodules:
- mcp.server.experimental.task_support.TaskSupport
- mcp.server.experimental.task_result_handler.TaskResultHandler
- mcp.server.experimental.request_context.Experimental
- mcp.server.experimental.task_scope (session scoping of task IDs)
"""

View file

@ -7,11 +7,15 @@ features within a request context, such as task-augmented request handling.
WARNING: These APIs are experimental and may change without notice.
"""
import warnings
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any
from typing import Any, overload
from typing_extensions import deprecated
from mcp.server.experimental.task_context import ServerTaskContext
from mcp.server.experimental.task_scope import scoped_task_id
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.session import ServerSession
from mcp.shared.exceptions import McpError
@ -29,6 +33,14 @@ from mcp.types import (
Tool,
)
EXPLICIT_TASK_ID_DEPRECATION = (
"Passing an explicit task_id to run_task is deprecated. A task created with an "
"explicit ID is not associated with the session that created it: any requestor "
"that presents the ID can read its status and result or cancel it, and it never "
"appears in tasks/list. Omit task_id to let the SDK generate an ID associated "
"with the creating session."
)
@dataclass
class Experimental:
@ -143,6 +155,25 @@ class Experimental:
return False
return True
@overload
async def run_task(
self,
work: Callable[[ServerTaskContext], Awaitable[Result]],
*,
task_id: None = None,
model_immediate_response: str | None = None,
) -> CreateTaskResult: ...
@overload
@deprecated(EXPLICIT_TASK_ID_DEPRECATION)
async def run_task(
self,
work: Callable[[ServerTaskContext], Awaitable[Result]],
*,
task_id: str,
model_immediate_response: str | None = None,
) -> CreateTaskResult: ...
async def run_task(
self,
work: Callable[[ServerTaskContext], Awaitable[Result]],
@ -167,9 +198,17 @@ class Experimental:
When work() returns a Result, the task is auto-completed with that result.
If work() raises an exception, the task is auto-failed.
Generated task IDs embed the session's task scope so that the default
task handlers only serve the task to the session that created it. An
explicitly provided `task_id` is used verbatim and is not associated
with the session, so any session can access it through the default
handlers; passing one is deprecated for that reason.
Args:
work: Async function that does the actual work
task_id: Optional task ID (generated if not provided)
task_id: Deprecated. Optional task ID, used verbatim and not
associated with the creating session. Omit it to let the SDK
generate one.
model_immediate_response: Optional string to include in _meta as
io.modelcontextprotocol/model-immediate-response
@ -196,6 +235,8 @@ class Experimental:
WARNING: This API is experimental and may change without notice.
"""
if task_id is not None:
warnings.warn(EXPLICIT_TASK_ID_DEPRECATION, DeprecationWarning, stacklevel=2)
if self._task_support is None:
raise RuntimeError("Task support not enabled. Call server.experimental.enable_tasks() first.")
if self._session is None:
@ -210,6 +251,11 @@ class Experimental:
# Access task_group via TaskSupport - raises if not in run() context
task_group = support.task_group
if task_id is None:
session_scope = self._session.experimental.task_session_scope
if session_scope is not None:
task_id = scoped_task_id(session_scope)
task = await support.store.create_task(self.task_metadata, task_id)
task_ctx = ServerTaskContext(

View file

@ -40,6 +40,12 @@ class ExperimentalServerSessionFeatures:
def __init__(self, session: "ServerSession") -> None:
self._session = session
# Opaque marker identifying this session for task scoping. Assigned by
# TaskSupport.configure_session(). Task IDs generated by run_task()
# embed it so the default task handlers can restrict task access to
# the session that created the task. None means tasks created on this
# session are not associated with it (e.g. stateless servers).
self.task_session_scope: str | None = None
async def get_task(self, task_id: str) -> types.GetTaskResult:
"""

View file

@ -46,6 +46,11 @@ class TaskResultHandler:
4. Blocks until task reaches terminal state
5. Returns the final result
Prefer `server.experimental.enable_tasks()`, whose default tasks/result
handler wraps `handle()` and only serves tasks created by the requesting
session. A custom handler that calls `handle()` directly is responsible
for deciding which requestors may access which tasks.
Usage:
# Create handler with store and queue
handler = TaskResultHandler(task_store, message_queue)
@ -55,9 +60,6 @@ class TaskResultHandler:
async def handle_task_result(req: GetTaskPayloadRequest) -> GetTaskPayloadResult:
ctx = server.request_context
return await handler.handle(req, ctx.session, ctx.request_id)
# Or use the convenience method
handler.register(server)
"""
def __init__(

View file

@ -0,0 +1,75 @@
"""
Session scoping for experimental task identifiers.
Task IDs generated by `run_task()` embed an opaque, per-session marker (the
"session scope") so that the default task handlers can tell which session
created a task. The default handlers for tasks/get, tasks/result, tasks/list,
and tasks/cancel only operate on tasks created by the requesting session.
Task IDs without a session scope (explicitly provided IDs, IDs created
directly through a TaskStore, or IDs created in stateless mode) have no known
creator. They can be used with tasks/get, tasks/result, and tasks/cancel from
any session - possession of the ID is what grants access - but they are never
included in tasks/list responses.
WARNING: These APIs are experimental and may change without notice.
"""
import re
from uuid import uuid4
__all__ = [
"new_session_scope",
"scoped_task_id",
"session_scope_of",
"task_in_session_scope",
"task_listable_in_session_scope",
]
# A scoped task ID has the form "<32 hex chars>:<uuid4>". Both halves must
# match exactly so that explicitly chosen task IDs are never mistaken for
# scoped ones. \Z rather than $ so a trailing newline cannot match.
_SCOPED_TASK_ID = re.compile(
r"\A(?P<scope>[0-9a-f]{32}):"
r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\Z"
)
def new_session_scope() -> str:
"""Create a new opaque session scope token."""
return uuid4().hex
def scoped_task_id(session_scope: str) -> str:
"""Generate a task ID associated with the given session scope."""
return f"{session_scope}:{uuid4()}"
def session_scope_of(task_id: str) -> str | None:
"""Return the session scope embedded in a task ID, or None if it has none."""
match = _SCOPED_TASK_ID.match(task_id)
return match.group("scope") if match else None
def task_in_session_scope(task_id: str, session_scope: str | None) -> bool:
"""Whether a task may be used by a requestor with the given session scope.
Used by tasks/get, tasks/result, and tasks/cancel. A task whose ID carries
no session scope has no known creator, so possession of the ID is what
grants access to it: it can be used from any session.
"""
embedded = session_scope_of(task_id)
return embedded is None or embedded == session_scope
def task_listable_in_session_scope(task_id: str, session_scope: str | None) -> bool:
"""Whether a task may be included in a tasks/list response for the given session scope.
Used by tasks/list. Listing is stricter than access by ID: a task is only
listed to the session that created it. Tasks with no session scope are
never listed because they have no known creator, and requestors with no
session scope are never shown any tasks because the server cannot tell
them apart.
"""
embedded = session_scope_of(task_id)
return embedded is not None and embedded == session_scope

View file

@ -13,6 +13,7 @@ import anyio
from anyio.abc import TaskGroup
from mcp.server.experimental.task_result_handler import TaskResultHandler
from mcp.server.experimental.task_scope import new_session_scope
from mcp.server.session import ServerSession
from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
@ -83,7 +84,7 @@ class TaskSupport:
finally:
self._task_group = None
def configure_session(self, session: ServerSession) -> None:
def configure_session(self, session: ServerSession, *, stateless: bool = False) -> None:
"""
Configure a session for task support.
@ -91,12 +92,22 @@ class TaskSupport:
responses to queued requests (elicitation, sampling) are routed
back to the waiting resolvers.
It also assigns the session a task session scope. Task IDs generated
by `run_task()` embed this scope, and the default task handlers only
operate on tasks created by the requesting session. Stateless sessions
are not assigned a scope: each request runs on a fresh session, so a
task created by one request could never be retrieved by a later one if
tasks were bound to the session that created them.
Called automatically by Server.run() for each new session.
Args:
session: The session to configure
stateless: Whether the session belongs to a stateless server run
"""
session.add_response_router(self.handler)
if not stateless and session.experimental.task_session_scope is None:
session.experimental.task_session_scope = new_session_scope()
@classmethod
def in_memory(cls) -> "TaskSupport":

View file

@ -9,6 +9,7 @@ import logging
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING
from mcp.server.experimental.task_scope import task_in_session_scope, task_listable_in_session_scope
from mcp.server.experimental.task_support import TaskSupport
from mcp.server.lowlevel.func_inspection import create_call_wrapper
from mcp.shared.exceptions import McpError
@ -31,6 +32,7 @@ from mcp.types import (
ServerResult,
ServerTasksCapability,
ServerTasksRequestsCapability,
Task,
TasksCallCapability,
TasksCancelCapability,
TasksListCapability,
@ -125,8 +127,38 @@ class ExperimentalHandlers:
return self._task_support
def _requestor_session_scope(self) -> str | None:
"""Return the task session scope of the session making the current request."""
return self._server.request_context.session.experimental.task_session_scope
def _require_task_in_requestor_scope(self, task_id: str) -> None:
"""Reject task IDs that belong to a different session.
Task IDs generated by `run_task()` embed the creating session's
scope. The default handlers treat a task created by another session
exactly like a task that does not exist, so a requestor cannot tell
whether such a task exists. Task IDs without an embedded scope are
accepted from any session.
Raises:
McpError: With INVALID_PARAMS if the task belongs to another session.
"""
if not task_in_session_scope(task_id, self._requestor_session_scope()):
raise McpError(
ErrorData(
code=INVALID_PARAMS,
message=f"Task not found: {task_id}",
)
)
def _register_default_task_handlers(self) -> None:
"""Register default handlers for task operations."""
"""Register default handlers for task operations.
Each default handler only operates on tasks created by the requesting
session (see `_require_task_in_requestor_scope`), and tasks/list only
returns the requesting session's own tasks (see
`task_listable_in_session_scope`).
"""
assert self._task_support is not None
support = self._task_support
@ -134,6 +166,7 @@ class ExperimentalHandlers:
if GetTaskRequest not in self._request_handlers:
async def _default_get_task(req: GetTaskRequest) -> ServerResult:
self._require_task_in_requestor_scope(req.params.taskId)
task = await support.store.get_task(req.params.taskId)
if task is None:
raise McpError(
@ -160,6 +193,7 @@ class ExperimentalHandlers:
if GetTaskPayloadRequest not in self._request_handlers:
async def _default_get_task_result(req: GetTaskPayloadRequest) -> ServerResult:
self._require_task_in_requestor_scope(req.params.taskId)
ctx = self._server.request_context
result = await support.handler.handle(req, ctx.session, ctx.request_id)
return ServerResult(result)
@ -170,9 +204,26 @@ class ExperimentalHandlers:
if ListTasksRequest not in self._request_handlers:
async def _default_list_tasks(req: ListTasksRequest) -> ServerResult:
cursor = req.params.cursor if req.params else None
tasks, next_cursor = await support.store.list_tasks(cursor)
return ServerResult(ListTasksResult(tasks=tasks, nextCursor=next_cursor))
requestor_scope = self._requestor_session_scope()
if requestor_scope is None:
# The server cannot tell this requestor apart from any
# other, so there are no tasks it can be shown.
return ServerResult(ListTasksResult(tasks=[]))
# Return every task that belongs to the requesting session in
# a single page. The store's pagination cursor is never sent
# to the requestor: it is derived from the unfiltered listing,
# so it could identify a task belonging to a different
# session. For the same reason the request's cursor is not
# forwarded to the store.
own_tasks: list[Task] = []
cursor: str | None = None
while True:
page, cursor = await support.store.list_tasks(cursor)
own_tasks.extend(
task for task in page if task_listable_in_session_scope(task.taskId, requestor_scope)
)
if cursor is None:
return ServerResult(ListTasksResult(tasks=own_tasks))
self._request_handlers[ListTasksRequest] = _default_list_tasks
@ -180,6 +231,7 @@ class ExperimentalHandlers:
if CancelTaskRequest not in self._request_handlers:
async def _default_cancel_task(req: CancelTaskRequest) -> ServerResult:
self._require_task_in_requestor_scope(req.params.taskId)
result = await cancel_task(support.store, req.params.taskId)
return ServerResult(result)

View file

@ -667,7 +667,7 @@ class Server(Generic[LifespanResultT, RequestT]):
# Configure task support for this session if enabled
task_support = self._experimental_handlers.task_support if self._experimental_handlers else None
if task_support is not None:
task_support.configure_session(session)
task_support.configure_session(session, stateless=stateless)
await stack.enter_async_context(task_support.run())
async with anyio.create_task_group() as tg: