ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-20 20:44:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Scope tasks to authorization context, not session (#3800)

Browse source 
* Scope tasks to authorization context, not session

Tasks were keyed by the transport-layer Mcp-Session-Id, which is
server-assigned and changes on reconnect — so clients lost access to
their running tasks after any connection interruption.

The MCP spec says tasks should be bound to authorization context, not
session.  This replaces session_id with task_scope (derived from
AccessToken.client_id, URL-encoded) in all task data Redis keys and
Docket task keys.  When no auth is configured, a "_" sentinel is used
and security comes from UUID task ID entropy per the spec.

Session ID is still used for transport-level concerns (notification
queues, subscriber registration) and is now stored in the
TaskContextSnapshot payload so background workers can still deliver
notifications.

Also extracts all the task context infrastructure (TaskContextInfo,
TaskContextSnapshot, snapshot loading, session/server registries) from
server/dependencies.py into a new server/tasks/context.py to keep the
DI module from sprawling further.

Closes #3758

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Rename _redis_key to _snapshot_redis_key

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Document in-process session registry as an optimization

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Tidy imports and docstrings

Hoist imports where safe, keep subscriptions/notifications deferred in
handlers.py since they pull in docket at module level. Sharpen docstrings
on keys.py and context.py so each module owns its lane. Clean up the
re-export block in dependencies.py.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix misleading comment on re-export block

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Tighten task scope: include sub claim, partition keyspaces

Addresses review feedback on #3800:

- Compose task scope from client_id and the JWT sub claim (when present)
  so fixed-OAuth deployments isolate per user, not just per client.
- Replace the "_" anonymous sentinel with a tagged keyspace partition.
  Docket keys are now auth:{enc_scope}:... or anon:..., and Redis keys
  use fastmcp:task:auth:{enc_scope}:... or fastmcp:task:anon:...,
  routed through a single task_redis_prefix() helper.
- get_task_scope() returns the raw scope (or None); encoding happens
  once at the keys.py boundary, collapsing the previous double-quote
  invariant.
- Drop the dormant fallback in notifications.py that routed
  input_required relays into the anon keyspace when task_scope was
  missing -- log and skip instead.
- Add comprehensive parser/encoder tests in test_task_keys.py covering
  round-trips, malformed keys, and adversarial scopes ("anon", "_", and
  scopes containing : / | %).
- Add cross-scope rejection tests: distinct client_ids, distinct sub
  claims under a shared client_id, and authenticated vs anonymous.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Guidry 2026-04-13 13:50:43 -04:00 • committed by William Easton
commit 08d4059f19
No known key found for this signature in database
11 changed files with 954 additions and 511 deletions
Download patch file Download diff file Expand all files Collapse all files

363
src/fastmcp/server/dependencies.py
View file

@ -10,14 +10,10 @@ from __future__ import annotations
import contextlib
import importlib.metadata
import inspect
import json
import logging
import weakref
from collections import OrderedDict
from collections.abc import AsyncGenerator, Callable
from contextlib import AsyncExitStack, asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache
from types import TracebackType
@ -45,12 +41,9 @@ from fastmcp.utilities.async_utils import (
)
from fastmcp.utilities.types import find_kwarg_by_type, is_class_member_of_type
_logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from docket import Docket
from docket.worker import Worker
from mcp.server.session import ServerSession
from fastmcp.server.context import Context
from fastmcp.server.server import FastMCP
@ -86,337 +79,30 @@ __all__ = [
]
# --- TaskContextInfo and get_task_context ---
@dataclass(frozen=True, slots=True)
class TaskContextInfo:
"""Information about the current background task context.
Returned by ``get_task_context()`` when running inside a Docket worker.
Contains identifiers needed to communicate with the MCP session.
"""
task_id: str
"""The MCP task ID (server-generated UUID)."""
session_id: str
"""The session ID that submitted this task."""
def get_task_context() -> TaskContextInfo | None:
"""Get the current task context if running inside a background task worker.
This function extracts task information from the Docket execution context.
Returns None if not running in a task context (e.g., foreground execution).
Returns:
TaskContextInfo with task_id and session_id, or None if not in a task.
"""
if not is_docket_available():
return None
from docket.dependencies import current_execution
try:
execution = current_execution.get()
# Parse the task key: {session_id}:{task_id}:{task_type}:{component}
from fastmcp.server.tasks.keys import parse_task_key
key_parts = parse_task_key(execution.key)
return TaskContextInfo(
task_id=key_parts["client_task_id"],
session_id=key_parts["session_id"],
)
except LookupError:
# Not in worker context
return None
except (ValueError, KeyError):
# Invalid task key format
return None
# --- Session registry for background task Context ---
_task_sessions: dict[str, weakref.ref[ServerSession]] = {}
def register_task_session(session_id: str, session: ServerSession) -> None:
"""Register a session for Context access in background tasks.
Called automatically when a task is submitted to Docket. The session is
stored as a weakref so it doesn't prevent garbage collection when the
client disconnects.
Args:
session_id: The session identifier
session: The ServerSession instance
"""
_task_sessions[session_id] = weakref.ref(session)
def get_task_session(session_id: str) -> ServerSession | None:
"""Get a registered session by ID if still alive.
Args:
session_id: The session identifier
Returns:
The ServerSession if found and alive, None otherwise
"""
ref = _task_sessions.get(session_id)
if ref is None:
return None
session = ref()
if session is None:
# Session was garbage collected, clean up entry
_task_sessions.pop(session_id, None)
return session
# --- ContextVars ---
# Task context lives in fastmcp.server.tasks.context; public symbols are
# re-exported here so existing imports from dependencies continue to work.
# _get_task_snapshot_sync and _load_task_snapshot_async are not re-exported
# but are used internally by get_access_token / get_http_request / get_server.
from fastmcp.server.tasks.context import (
TaskContextInfo,
TaskContextSnapshot,
_get_task_snapshot_sync,
_load_task_snapshot_async,
get_task_context,
get_task_server,
get_task_session,
register_task_server,
register_task_session,
)
_current_server: ContextVar[weakref.ref[FastMCP] | None] = ContextVar(
"server", default=None
)
# --- Background task server map ---
# Maps task_id → server weakref so background workers can resolve the correct
# server for mounted-child tasks. Follows the same pattern as _task_sessions.
# Populated in submit_to_docket() where the child server is in context;
# consulted in get_server() when running inside a Docket worker.
_task_server_map: OrderedDict[str, weakref.ref[FastMCP]] = OrderedDict()
_TASK_SERVER_MAP_MAX_SIZE = 10_000
def register_task_server(task_id: str, server: FastMCP) -> None:
"""Register the server for a background task.
Called at task-submission time (inside the child server's call_tool
context) so that background workers can resolve CurrentFastMCP() and
ctx.fastmcp to the child server for mounted tasks.
The map is bounded to avoid unbounded growth in long-lived servers.
Evicted entries fall back to the ContextVar (parent server).
"""
_task_server_map[task_id] = weakref.ref(server)
while len(_task_server_map) > _TASK_SERVER_MAP_MAX_SIZE:
_task_server_map.popitem(last=False)
_current_docket: ContextVar[Docket | None] = ContextVar("docket", default=None)
_current_worker: ContextVar[Worker | None] = ContextVar("worker", default=None)
# --- Unified task context snapshot ---
@dataclass(frozen=True, slots=True)
class TaskContextSnapshot:
"""All context data snapshotted at task-submission time.
Stored as a single Redis key per task, restored once in the worker.
"""
access_token_json: str | None = None
http_headers: dict[str, str] | None = None
origin_request_id: str | None = None
@classmethod
def capture(cls) -> TaskContextSnapshot:
"""Capture current context for background task execution."""
access_token = get_access_token()
ctx = get_context()
request_context = ctx.request_context
return cls(
access_token_json=(
access_token.model_dump_json() if access_token else None
),
http_headers=get_http_headers(include_all=True) or None,
origin_request_id=(
str(request_context.request_id) if request_context is not None else None
),
)
@classmethod
def from_json(cls, raw: str | bytes) -> TaskContextSnapshot:
"""Deserialize from JSON stored in Redis."""
if isinstance(raw, bytes):
raw = raw.decode()
parsed = json.loads(raw)
headers = parsed.get("http_headers")
if isinstance(headers, dict):
headers = {str(k).lower(): str(v) for k, v in headers.items()}
return cls(
access_token_json=parsed.get("access_token_json"),
http_headers=headers,
origin_request_id=parsed.get("origin_request_id"),
)
def to_json(self) -> str:
"""Serialize to JSON for Redis storage."""
return json.dumps(
{
"access_token_json": self.access_token_json,
"http_headers": self.http_headers,
"origin_request_id": self.origin_request_id,
}
)
async def save(
self,
docket: Docket,
session_id: str,
task_id: str,
ttl_seconds: int,
) -> None:
"""Store this snapshot as a single Redis key."""
key = docket.key(f"fastmcp:task:{session_id}:{task_id}:snapshot")
async with docket.redis() as redis:
await redis.set(key, self.to_json(), ex=ttl_seconds)
# Cache keyed by task_id so stale entries from previous tasks in the same
# asyncio context are automatically ignored (Docket workers may reuse contexts).
_task_snapshot: ContextVar[tuple[str, TaskContextSnapshot] | None] = ContextVar(
"task_snapshot", default=None
)
def _set_cached_snapshot(task_id: str, snapshot: TaskContextSnapshot) -> None:
"""Cache a snapshot keyed by task_id."""
_task_snapshot.set((task_id, snapshot))
def _get_cached_snapshot(task_id: str) -> TaskContextSnapshot | None:
"""Get cached snapshot if it belongs to this task."""
cached = _task_snapshot.get()
if cached is not None:
cached_task_id, snapshot = cached
if cached_task_id == task_id:
return snapshot
return None
def _redis_key(session_id: str, task_id: str) -> str:
"""Build the Redis key suffix for a task snapshot."""
return f"fastmcp:task:{session_id}:{task_id}:snapshot"
async def _load_task_snapshot_async(
session_id: str, task_id: str
) -> TaskContextSnapshot | None:
"""Load task context snapshot from Redis (async) and cache it.
Idempotent — returns the cached value if already loaded for this task.
"""
cached = _get_cached_snapshot(task_id)
if cached is not None:
return cached
try:
docket = get_server()._docket
except RuntimeError:
docket = None
if docket is None:
docket = _current_docket.get()
if docket is None:
return None
try:
async with docket.redis() as redis:
raw = await redis.get(docket.key(_redis_key(session_id, task_id)))
if raw is None:
return None
snapshot = TaskContextSnapshot.from_json(raw)
_set_cached_snapshot(task_id, snapshot)
return snapshot
except (OSError, json.JSONDecodeError, KeyError, ValueError):
_logger.warning(
"Failed to load task snapshot for %s:%s",
session_id,
task_id,
exc_info=True,
)
return None
def _get_task_snapshot_sync() -> TaskContextSnapshot | None:
"""Get the task snapshot using only sync operations.
Fallback chain:
1. ContextVar cache (keyed by task_id, set by async or sync loaders)
2. Sync Redis GET (works for both memory:// and real Redis)
"""
task_info = get_task_context()
if task_info is None:
return None
cached = _get_cached_snapshot(task_info.task_id)
if cached is not None:
return cached
return _load_task_snapshot_sync(task_info.session_id, task_info.task_id)
def _load_task_snapshot_sync(
session_id: str, task_id: str
) -> TaskContextSnapshot | None:
"""Load snapshot via sync Redis.
For memory:// backends (fakeredis), shares the same FakeServer instance
that Docket uses so data is accessible. For real Redis, creates a standard
sync connection.
"""
try:
from docket.dependencies import current_docket as _docket_cv
docket = _docket_cv.get()
except (LookupError, ImportError):
return None
if docket is None:
return None
try:
sync_redis = _get_sync_redis(docket.url)
raw = sync_redis.get(docket.key(_redis_key(session_id, task_id)))
if raw is None:
return None
snapshot = TaskContextSnapshot.from_json(raw)
_set_cached_snapshot(task_id, snapshot)
return snapshot
except (OSError, json.JSONDecodeError, KeyError, ValueError, ImportError):
_logger.warning(
"Failed to load task snapshot via sync Redis for %s:%s",
session_id,
task_id,
exc_info=True,
)
return None
def _get_sync_redis(url: str) -> Any:
"""Get a sync Redis client that shares the same backend as Docket.
For memory:// URLs, connects to the same fakeredis FakeServer instance
so data written by the async Docket client is visible. For real Redis
URLs, creates a standard sync connection.
"""
from docket._redis import get_memory_server
server = get_memory_server(url)
if server is not None:
from fakeredis import FakeRedis
return FakeRedis(server=server)
from redis import Redis
return Redis.from_url(url)
# --- Docket availability check ---
_DOCKET_AVAILABLE: bool | None = None
@ -662,13 +348,9 @@ def get_server() -> FastMCP:
# This handles mounted-child tasks where _current_server is the parent.
task_info = get_task_context()
if task_info is not None:
ref = _task_server_map.get(task_info.task_id)
if ref is not None:
server = ref()
if server is not None:
return server
# Server was garbage collected, clean up
_task_server_map.pop(task_info.task_id, None)
task_server = get_task_server(task_info.task_id)
if task_server is not None:
return task_server
server_ref = _current_server.get()
if server_ref is None:
@ -1073,15 +755,20 @@ class _CurrentContext(Dependency["Context"]):
# Check if we're in a Docket worker context
task_info = get_task_context()
if task_info is not None:
session = get_task_session(task_info.session_id)
server = get_server()
# Load unified snapshot (sets _task_snapshot ContextVar)
snapshot = await _load_task_snapshot_async(
task_info.session_id, task_info.task_id
task_info.task_scope, task_info.task_id
)
origin_request_id = snapshot.origin_request_id if snapshot else None
# Session ID is stored in the snapshot for notification delivery
snapshot_session_id = snapshot.session_id if snapshot else None
session = (
get_task_session(snapshot_session_id) if snapshot_session_id else None
)
ctx = Context(
fastmcp=server,
session=session,

389
src/fastmcp/server/tasks/context.py Normal file
View file

@ -0,0 +1,389 @@
"""Task context and scoping for background task execution.
Determines authorization scope (``get_task_scope``), manages the context
snapshot that is captured at task submission and restored in workers
(``TaskContextSnapshot``), and maintains in-process registries for live
sessions and servers.
"""
from __future__ import annotations
import json
import logging
import weakref
from collections import OrderedDict
from contextvars import ContextVar
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from fastmcp.server.tasks.keys import parse_task_key, task_redis_prefix
if TYPE_CHECKING:
from docket import Docket
from mcp.server.session import ServerSession
from fastmcp.server.server import FastMCP
_logger = logging.getLogger(__name__)
def get_task_scope() -> str | None:
"""Get the authorization scope for task isolation.
Returns the raw scope identifier for the current access token, or
``None`` when no auth context is present (anonymous tasks).
The scope is composed as ``client_id|sub`` when the token carries a
``sub`` claim — necessary for fixed-OAuth servers where ``client_id`` is
shared across all users — and falls back to ``client_id`` alone for
DCR/CIMD flows where the client identity is already per-user.
Encoding for Redis/Docket keys happens at the boundary in ``keys.py``;
this function returns the raw value.
"""
from fastmcp.server.dependencies import get_access_token
token = get_access_token()
if token is None:
return None
sub = token.claims.get("sub") if token.claims else None
if sub:
return f"{token.client_id}|{sub}"
return token.client_id
@dataclass(frozen=True, slots=True)
class TaskContextInfo:
"""Information about the current background task context.
Returned by ``get_task_context()`` when running inside a Docket worker.
Contains identifiers needed to communicate with the MCP session.
"""
task_id: str
"""The MCP task ID (server-generated UUID)."""
task_scope: str | None
"""The authorization scope that owns this task, or ``None`` if anonymous."""
def get_task_context() -> TaskContextInfo | None:
"""Get the current task context if running inside a background task worker.
This function extracts task information from the Docket execution context.
Returns None if not running in a task context (e.g., foreground execution).
Returns:
TaskContextInfo with task_id and task_scope, or None if not in a task.
"""
from fastmcp.server.dependencies import is_docket_available
if not is_docket_available():
return None
from docket.dependencies import current_execution
try:
execution = current_execution.get()
key_parts = parse_task_key(execution.key)
return TaskContextInfo(
task_id=key_parts["client_task_id"],
task_scope=key_parts["task_scope"],
)
except LookupError:
return None
except (ValueError, KeyError):
return None
@dataclass(frozen=True, slots=True)
class TaskContextSnapshot:
"""All context data snapshotted at task-submission time.
Stored as a single Redis key per task, restored once in the worker.
"""
access_token_json: str | None = None
http_headers: dict[str, str] | None = None
origin_request_id: str | None = None
session_id: str | None = None
@classmethod
def capture(cls) -> TaskContextSnapshot:
"""Capture current context for background task execution."""
from fastmcp.server.dependencies import (
get_access_token,
get_context,
get_http_headers,
)
access_token = get_access_token()
ctx = get_context()
request_context = ctx.request_context
try:
session_id = ctx.session_id
except RuntimeError:
session_id = None
return cls(
access_token_json=(
access_token.model_dump_json() if access_token else None
),
http_headers=get_http_headers(include_all=True) or None,
origin_request_id=(
str(request_context.request_id) if request_context is not None else None
),
session_id=session_id,
)
@classmethod
def from_json(cls, raw: str | bytes) -> TaskContextSnapshot:
"""Deserialize from JSON stored in Redis."""
if isinstance(raw, bytes):
raw = raw.decode()
parsed = json.loads(raw)
headers = parsed.get("http_headers")
if isinstance(headers, dict):
headers = {str(k).lower(): str(v) for k, v in headers.items()}
return cls(
access_token_json=parsed.get("access_token_json"),
http_headers=headers,
origin_request_id=parsed.get("origin_request_id"),
session_id=parsed.get("session_id"),
)
def to_json(self) -> str:
"""Serialize to JSON for Redis storage."""
return json.dumps(
{
"access_token_json": self.access_token_json,
"http_headers": self.http_headers,
"origin_request_id": self.origin_request_id,
"session_id": self.session_id,
}
)
async def save(
self,
docket: Docket,
task_scope: str | None,
task_id: str,
ttl_seconds: int,
) -> None:
"""Store this snapshot as a single Redis key."""
key = docket.key(_snapshot_redis_key(task_scope, task_id))
async with docket.redis() as redis:
await redis.set(key, self.to_json(), ex=ttl_seconds)
# Cache keyed by task_id so stale entries from previous tasks in the same
# asyncio context are automatically ignored (Docket workers may reuse contexts).
_task_snapshot: ContextVar[tuple[str, TaskContextSnapshot] | None] = ContextVar(
"task_snapshot", default=None
)
def _set_cached_snapshot(task_id: str, snapshot: TaskContextSnapshot) -> None:
"""Cache a snapshot keyed by task_id."""
_task_snapshot.set((task_id, snapshot))
def _get_cached_snapshot(task_id: str) -> TaskContextSnapshot | None:
"""Get cached snapshot if it belongs to this task."""
cached = _task_snapshot.get()
if cached is not None:
cached_task_id, snapshot = cached
if cached_task_id == task_id:
return snapshot
return None
def _snapshot_redis_key(task_scope: str | None, task_id: str) -> str:
"""Build the Redis key suffix for a task snapshot."""
return f"{task_redis_prefix(task_scope)}:{task_id}:snapshot"
async def _load_task_snapshot_async(
task_scope: str | None, task_id: str
) -> TaskContextSnapshot | None:
"""Load task context snapshot from Redis (async) and cache it.
Idempotent — returns the cached value if already loaded for this task.
"""
cached = _get_cached_snapshot(task_id)
if cached is not None:
return cached
from fastmcp.server.dependencies import _current_docket, get_server
try:
docket = get_server()._docket
except RuntimeError:
docket = None
if docket is None:
docket = _current_docket.get()
if docket is None:
return None
try:
async with docket.redis() as redis:
raw = await redis.get(docket.key(_snapshot_redis_key(task_scope, task_id)))
if raw is None:
return None
snapshot = TaskContextSnapshot.from_json(raw)
_set_cached_snapshot(task_id, snapshot)
return snapshot
except (OSError, json.JSONDecodeError, KeyError, ValueError):
_logger.warning(
"Failed to load task snapshot for %s:%s",
task_scope,
task_id,
exc_info=True,
)
return None
def get_task_session_id() -> str | None:
"""Get the session_id for the current background task, if available.
Loads the task snapshot (from cache or Redis) and returns the session_id
that was captured at task submission time. Returns None if not in a task
context or if the snapshot isn't available.
"""
snapshot = _get_task_snapshot_sync()
return snapshot.session_id if snapshot else None
def _get_task_snapshot_sync() -> TaskContextSnapshot | None:
"""Get the task snapshot using only sync operations.
Fallback chain:
1. ContextVar cache (keyed by task_id, set by async or sync loaders)
2. Sync Redis GET (works for both memory:// and real Redis)
"""
task_info = get_task_context()
if task_info is None:
return None
cached = _get_cached_snapshot(task_info.task_id)
if cached is not None:
return cached
return _load_task_snapshot_sync(task_info.task_scope, task_info.task_id)
def _load_task_snapshot_sync(
task_scope: str | None, task_id: str
) -> TaskContextSnapshot | None:
"""Load snapshot via sync Redis.
For memory:// backends (fakeredis), shares the same FakeServer instance
that Docket uses so data is accessible. For real Redis, creates a standard
sync connection.
"""
try:
from docket.dependencies import current_docket as _docket_cv
docket = _docket_cv.get()
except (LookupError, ImportError):
return None
if docket is None:
return None
try:
sync_redis = _get_sync_redis(docket.url)
raw = sync_redis.get(docket.key(_snapshot_redis_key(task_scope, task_id)))
if raw is None:
return None
snapshot = TaskContextSnapshot.from_json(raw)
_set_cached_snapshot(task_id, snapshot)
return snapshot
except (OSError, json.JSONDecodeError, KeyError, ValueError, ImportError):
_logger.warning(
"Failed to load task snapshot via sync Redis for %s:%s",
task_scope,
task_id,
exc_info=True,
)
return None
def _get_sync_redis(url: str) -> Any:
"""Get a sync Redis client that shares the same backend as Docket.
For memory:// URLs, connects to the same fakeredis FakeServer instance
so data written by the async Docket client is visible. For real Redis
URLs, creates a standard sync connection.
"""
from docket._redis import get_memory_server
server = get_memory_server(url)
if server is not None:
from fakeredis import FakeRedis
return FakeRedis(server=server)
from redis import Redis
return Redis.from_url(url)
# In-process optimization: when the Docket worker runs in the same process as
# the MCP server, we can hand background tasks a live ServerSession so they can
# call session methods directly (e.g. send_notification). In distributed
# deployments where workers are separate processes, these registries will be
# empty and the worker's Context will have session=None — that's fine, because
# elicitation and notifications have Redis-based fallbacks that work across
# process boundaries (see notifications.py and elicitation.py).
_task_sessions: dict[str, weakref.ref[ServerSession]] = {}
def register_task_session(session_id: str, session: ServerSession) -> None:
"""Register a session for in-process background task access.
Called automatically when a task is submitted to Docket. The session is
stored as a weakref so it doesn't prevent garbage collection when the
client disconnects.
"""
_task_sessions[session_id] = weakref.ref(session)
def get_task_session(session_id: str) -> ServerSession | None:
"""Get a registered session by ID if still alive.
Returns None in distributed workers where the session lives in another
process — callers must handle this gracefully.
"""
ref = _task_sessions.get(session_id)
if ref is None:
return None
session = ref()
if session is None:
_task_sessions.pop(session_id, None)
return session
_task_server_map: OrderedDict[str, weakref.ref[FastMCP]] = OrderedDict()
_TASK_SERVER_MAP_MAX_SIZE = 10_000
def register_task_server(task_id: str, server: FastMCP) -> None:
"""Register the server for a background task.
Called at task-submission time so that background workers can resolve
the correct (child) server for mounted tasks.
"""
_task_server_map[task_id] = weakref.ref(server)
while len(_task_server_map) > _TASK_SERVER_MAP_MAX_SIZE:
_task_server_map.popitem(last=False)
def get_task_server(task_id: str) -> FastMCP | None:
"""Get the registered server for a background task, if still alive."""
ref = _task_server_map.get(task_id)
if ref is None:
return None
server = ref()
if server is None:
_task_server_map.pop(task_id, None)
return server

66
src/fastmcp/server/tasks/elicitation.py
View file

@ -24,21 +24,26 @@ from typing import TYPE_CHECKING, Any, cast
import mcp.types
from mcp import ServerSession
from fastmcp.server.tasks.context import get_task_context, get_task_session_id
from fastmcp.server.tasks.keys import task_redis_prefix
from fastmcp.server.tasks.notifications import push_notification
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from fastmcp.server.server import FastMCP
# Redis key patterns for task elicitation state
ELICIT_REQUEST_KEY = "fastmcp:task:{session_id}:{task_id}:elicit:request"
ELICIT_RESPONSE_KEY = "fastmcp:task:{session_id}:{task_id}:elicit:response"
ELICIT_STATUS_KEY = "fastmcp:task:{session_id}:{task_id}:elicit:status"
# TTL for elicitation state (1 hour)
ELICIT_TTL_SECONDS = 3600
def _elicit_keys(task_scope: str | None, task_id: str) -> tuple[str, str, str]:
"""Build (request, response, status) Redis keys for a task's elicitation."""
prefix = f"{task_redis_prefix(task_scope)}:{task_id}:elicit"
return f"{prefix}:request", f"{prefix}:response", f"{prefix}:status"
async def elicit_for_task(
task_id: str,
session: ServerSession | None,
@ -75,26 +80,22 @@ async def elicit_for_task(
# Generate a unique request ID for this elicitation
request_id = str(uuid.uuid4())
# Get session ID from task context (authoritative source for background tasks)
# This is extracted from the Docket execution key: {session_id}:{task_id}:...
from fastmcp.server.dependencies import get_task_context
task_context = get_task_context()
if task_context is not None:
session_id = task_context.session_id
task_scope = task_context.task_scope
# Prefer the live session's cached ID (always available in-process),
# fall back to the snapshot for distributed workers.
session_id = (
getattr(session, "_fastmcp_state_prefix", None) or get_task_session_id()
)
else:
# Fallback: try to get from session attribute (shouldn't happen in background)
session_id = getattr(session, "_fastmcp_state_prefix", None)
if session_id is None:
raise RuntimeError(
"Cannot determine session_id for elicitation. "
"This typically means elicit_for_task() was called outside a Docket worker context."
)
raise RuntimeError(
"Cannot determine task scope for elicitation. "
"This typically means elicit_for_task() was called outside a Docket worker context."
)
# Store elicitation request in Redis
request_key = ELICIT_REQUEST_KEY.format(session_id=session_id, task_id=task_id)
response_key = ELICIT_RESPONSE_KEY.format(session_id=session_id, task_id=task_id)
status_key = ELICIT_STATUS_KEY.format(session_id=session_id, task_id=task_id)
request_key, response_key, status_key = _elicit_keys(task_scope, task_id)
elicit_request = {
"request_id": request_id,
@ -138,6 +139,7 @@ async def elicit_for_task(
"taskId": task_id,
"status": "input_required",
"statusMessage": message,
"task_scope": task_scope,
"elicitation": {
"requestId": request_id,
"message": message,
@ -147,9 +149,12 @@ async def elicit_for_task(
},
}
# Push notification to Redis queue (works from any process)
# Server's subscriber loop will forward to client
from fastmcp.server.tasks.notifications import push_notification
if session_id is None:
logger.warning(
"No session_id available for task %s, cannot deliver elicitation notification",
task_id,
)
return mcp.types.ElicitResult(action="cancel", content=None)
try:
await push_notification(session_id, notification_dict, docket)
@ -233,7 +238,7 @@ async def elicit_for_task(
async def relay_elicitation(
session: ServerSession,
session_id: str,
task_scope: str | None,
task_id: str,
elicitation: dict[str, Any],
fastmcp: FastMCP,
@ -247,7 +252,7 @@ async def relay_elicitation(
Args:
session: MCP ServerSession
session_id: Session identifier
task_scope: Authorization scope for Redis key construction
task_id: Background task ID
elicitation: Elicitation metadata (message, requestedSchema)
fastmcp: FastMCP server instance
@ -259,7 +264,7 @@ async def relay_elicitation(
)
await handle_task_input(
task_id=task_id,
session_id=session_id,
task_scope=task_scope,
action=result.action,
content=result.content,
fastmcp=fastmcp,
@ -274,7 +279,7 @@ async def relay_elicitation(
# Push a cancel response so the worker's BLPOP doesn't block forever
success = await handle_task_input(
task_id=task_id,
session_id=session_id,
task_scope=task_scope,
action="cancel",
content=None,
fastmcp=fastmcp,
@ -289,7 +294,7 @@ async def relay_elicitation(
async def handle_task_input(
task_id: str,
session_id: str,
task_scope: str | None,
action: str,
content: dict[str, Any] | None,
fastmcp: FastMCP,
@ -301,7 +306,7 @@ async def handle_task_input(
Args:
task_id: The background task ID
session_id: The MCP session ID
task_scope: Authorization scope for Redis key construction
action: The elicitation action ("accept", "decline", "cancel")
content: The response content (for "accept" action)
fastmcp: The FastMCP server instance
@ -313,8 +318,7 @@ async def handle_task_input(
if docket is None:
return False
response_key = ELICIT_RESPONSE_KEY.format(session_id=session_id, task_id=task_id)
status_key = ELICIT_STATUS_KEY.format(session_id=session_id, task_id=task_id)
_, response_key, status_key = _elicit_keys(task_scope, task_id)
response = {
"action": action,

86
src/fastmcp/server/tasks/handlers.py
View file

@ -15,13 +15,17 @@ from mcp.shared.exceptions import McpError
from mcp.types import INTERNAL_ERROR, ErrorData
from fastmcp.server.dependencies import (
TaskContextSnapshot,
_current_docket,
get_context,
register_task_server,
)
from fastmcp.server.tasks.config import TaskMeta
from fastmcp.server.tasks.keys import build_task_key
from fastmcp.server.tasks.context import (
TaskContextSnapshot,
get_task_scope,
register_task_server,
register_task_session,
)
from fastmcp.server.tasks.keys import build_task_key, task_redis_prefix
from fastmcp.utilities.logging import get_logger
if TYPE_CHECKING:
@ -69,12 +73,16 @@ async def submit_to_docket(
# Record creation timestamp per SEP-1686 final spec (line 430)
created_at = datetime.now(timezone.utc)
# Get session ID - use "internal" for programmatic calls without MCP session
ctx = get_context()
# Authorization scope for task isolation (auth identity, or None for anonymous)
task_scope = get_task_scope()
# Transport session ID for notification delivery
try:
session_id = ctx.session_id
except RuntimeError:
session_id = "internal"
session_id = None
# Try the server's own Docket first; fall back to the ContextVar for
# mounted children (whose parent server owns the Docket instance).
@ -94,7 +102,7 @@ async def submit_to_docket(
register_task_server(server_task_id, ctx.fastmcp)
# Build full task key with embedded metadata
task_key = build_task_key(session_id, server_task_id, task_type, key)
task_key = build_task_key(task_scope, server_task_id, task_type, key)
# Determine TTL: use task_meta.ttl if provided, else docket default
if task_meta is not None and task_meta.ttl is not None:
@ -104,16 +112,14 @@ async def submit_to_docket(
ttl_seconds = int(ttl_ms / 1000) + TASK_MAPPING_TTL_BUFFER_SECONDS
# Store task metadata in Redis for protocol handlers
task_meta_key = docket.key(f"fastmcp:task:{session_id}:{server_task_id}")
created_at_key = docket.key(
f"fastmcp:task:{session_id}:{server_task_id}:created_at"
)
poll_interval_key = docket.key(
f"fastmcp:task:{session_id}:{server_task_id}:poll_interval"
)
prefix = task_redis_prefix(task_scope)
task_meta_key = docket.key(f"{prefix}:{server_task_id}")
created_at_key = docket.key(f"{prefix}:{server_task_id}:created_at")
poll_interval_key = docket.key(f"{prefix}:{server_task_id}:poll_interval")
poll_interval_ms = int(component.task_config.poll_interval.total_seconds() * 1000)
# Snapshot all context (access token, headers, origin request ID) as a single key
# Snapshot all context (access token, headers, origin request ID,
# and session_id for notification delivery in background workers)
snapshot = TaskContextSnapshot.capture()
async with docket.redis() as redis:
@ -121,14 +127,12 @@ async def submit_to_docket(
await redis.set(created_at_key, created_at.isoformat(), ex=ttl_seconds)
await redis.set(poll_interval_key, str(poll_interval_ms), ex=ttl_seconds)
await snapshot.save(docket, session_id, server_task_id, ttl_seconds)
await snapshot.save(docket, task_scope, server_task_id, ttl_seconds)
# Register session for Context access in background workers (SEP-1686)
# This enables elicitation/sampling from background tasks via weakref
# Skip for "internal" sessions (programmatic calls without MCP session)
if session_id != "internal":
from fastmcp.server.dependencies import register_task_session
# Skip when there is no session (programmatic calls without MCP session)
if session_id is not None:
register_task_session(session_id, ctx.session)
# Send an initial tasks/status notification before queueing.
@ -160,7 +164,7 @@ async def submit_to_docket(
# Queue function to Docket by key (result storage via execution_ttl)
# Use component.add_to_docket() which handles calling conventions
# `fn_key` is the function lookup key (e.g., "child_multiply")
# `task_key` is the task result key (e.g., "fastmcp:task:{session}:{task_id}:tool:child_multiply")
# `task_key` is the task result key (e.g., "fastmcp:task:{task_scope}:{task_id}:tool:child_multiply")
# Resources don't take arguments; tools/prompts/templates always pass arguments (even if None/empty)
if task_type == "resource":
await component.add_to_docket(docket, fn_key=key, task_key=task_key) # type: ignore[call-arg] # ty:ignore[missing-argument]
@ -168,9 +172,10 @@ async def submit_to_docket(
await component.add_to_docket(docket, arguments, fn_key=key, task_key=task_key) # type: ignore[call-arg] # ty:ignore[invalid-argument-type, too-many-positional-arguments]
# Spawn subscription task to send status notifications (SEP-1686 optional feature)
# Start subscription in session's task group (persists for connection lifetime)
# Deferred: subscriptions and notifications depend on docket at import time
from fastmcp.server.tasks.subscriptions import subscribe_to_task_updates
# Start subscription in session's task group (persists for connection lifetime)
if hasattr(ctx.session, "_subscription_task_group"):
tg = ctx.session._subscription_task_group
if tg:
@ -183,33 +188,34 @@ async def submit_to_docket(
poll_interval_ms,
)
# Start notification subscriber for distributed elicitation (idempotent)
# This enables ctx.elicit() to work when workers run in separate processes
# Subscriber forwards notifications from Redis queue to client session
# Deferred: notifications depends on docket at import time
from fastmcp.server.tasks.notifications import (
ensure_subscriber_running,
stop_subscriber,
)
try:
await ensure_subscriber_running(session_id, ctx.session, docket, ctx.fastmcp)
if session_id is not None:
try:
await ensure_subscriber_running(
session_id, ctx.session, docket, ctx.fastmcp
)
# Register cleanup callback on session exit (once per session)
# This ensures subscriber is stopped when the session disconnects
if (
hasattr(ctx.session, "_exit_stack")
and ctx.session._exit_stack is not None
and not getattr(ctx.session, "_notification_cleanup_registered", False)
):
# Register cleanup callback on session exit (once per session)
# This ensures subscriber is stopped when the session disconnects
if (
hasattr(ctx.session, "_exit_stack")
and ctx.session._exit_stack is not None
and not getattr(ctx.session, "_notification_cleanup_registered", False)
):
async def _cleanup_subscriber() -> None:
await stop_subscriber(session_id)
async def _cleanup_subscriber() -> None:
await stop_subscriber(session_id) # type: ignore[arg-type]
ctx.session._exit_stack.push_async_callback(_cleanup_subscriber)
ctx.session._notification_cleanup_registered = True # type: ignore[attr-defined] # ty:ignore[unresolved-attribute]
except Exception as e:
# Non-fatal: elicitation will still work via polling fallback
logger.debug("Failed to start notification subscriber: %s", e)
ctx.session._exit_stack.push_async_callback(_cleanup_subscriber)
ctx.session._notification_cleanup_registered = True # type: ignore[attr-defined] # ty:ignore[unresolved-attribute]
except Exception as e:
# Non-fatal: elicitation will still work via polling fallback
logger.debug("Failed to start notification subscriber: %s", e)
# Return CreateTaskResult with proper Task object
# Tasks MUST begin in "working" status per SEP-1686 final spec (line 381)

142
src/fastmcp/server/tasks/keys.py
View file

@ -1,31 +1,56 @@
"""Task key management for SEP-1686 background tasks.
"""Docket and Redis key encoding for background tasks.
Task keys encode security scoping and metadata in the Docket key format:
`{session_id}:{client_task_id}:{task_type}:{component_identifier}`
The compound Docket task key embeds the auth boundary so that the parser can
reject cross-scope access without consulting Redis. Authenticated and
anonymous tasks live in disjoint keyspaces:
This format provides:
- Session-based security scoping (prevents cross-session access)
- Task type identification (tool/prompt/resource)
- Component identification (name or URI for result conversion)
auth:{enc_scope}:{client_task_id}:{task_type}:{enc_identifier}
anon:{client_task_id}:{task_type}:{enc_identifier}
The same `auth/anon` partition is used for the per-task Redis prefix
(``fastmcp:task:auth:{enc_scope}`` vs ``fastmcp:task:anon``) — see
``task_redis_prefix``.
``task_scope`` is the raw scope identifier (typically derived from
``client_id`` or ``client_id|sub``); encoding happens once, at the boundary,
in this module.
"""
from typing import TypedDict
from urllib.parse import quote, unquote
class TaskKeyParts(TypedDict):
"""Decoded segments of a Docket task key.
``task_scope`` is ``None`` for anonymous tasks, the raw scope string
otherwise.
"""
task_scope: str | None
client_task_id: str
task_type: str
component_identifier: str
_AUTH_TAG = "auth"
_ANON_TAG = "anon"
_VALID_TAGS = (_AUTH_TAG, _ANON_TAG)
def build_task_key(
session_id: str,
task_scope: str | None,
client_task_id: str,
task_type: str,
component_identifier: str,
) -> str:
"""Build Docket task key with embedded metadata.
Format: `{session_id}:{client_task_id}:{task_type}:{component_identifier}`
The component_identifier is URI-encoded to handle special characters (colons, slashes, etc.).
When ``task_scope`` is ``None`` the task is anonymous and lives in the
``anon`` keyspace. Otherwise it lives under ``auth:{enc_scope}``.
Args:
session_id: Session ID for security scoping
task_scope: Raw authorization scope, or ``None`` for anonymous tasks
client_task_id: Client-provided task ID
task_type: Type of task ("tool", "prompt", "resource")
component_identifier: Tool name, prompt name, or resource URI
@ -34,44 +59,78 @@ def build_task_key(
Encoded task key for Docket
Examples:
>>> build_task_key("session123", "task456", "tool", "my_tool")
'session123:task456:tool:my_tool'
>>> build_task_key("client-a", "task456", "tool", "my_tool")
'auth:client-a:task456:tool:my_tool'
>>> build_task_key("session123", "task456", "resource", "file://data.txt")
'session123:task456:resource:file%3A%2F%2Fdata.txt'
>>> build_task_key(None, "task456", "tool", "my_tool")
'anon:task456:tool:my_tool'
>>> build_task_key("client-a", "task456", "resource", "file://data.txt")
'auth:client-a:task456:resource:file%3A%2F%2Fdata.txt'
"""
encoded_identifier = quote(component_identifier, safe="")
return f"{session_id}:{client_task_id}:{task_type}:{encoded_identifier}"
if task_scope is None:
return f"{_ANON_TAG}:{client_task_id}:{task_type}:{encoded_identifier}"
encoded_scope = quote(task_scope, safe="")
return (
f"{_AUTH_TAG}:{encoded_scope}:{client_task_id}:{task_type}:{encoded_identifier}"
)
def parse_task_key(task_key: str) -> dict[str, str]:
def parse_task_key(task_key: str) -> TaskKeyParts:
"""Parse Docket task key to extract metadata.
Args:
task_key: Encoded task key from Docket
Returns:
Dict with keys: session_id, client_task_id, task_type, component_identifier
Dict with keys: ``task_scope`` (``str | None``), ``client_task_id``,
``task_type``, ``component_identifier``.
Raises:
ValueError: If the key has an unrecognized tag or wrong segment count.
Examples:
>>> parse_task_key("session123:task456:tool:my_tool")
`{'session_id': 'session123', 'client_task_id': 'task456', 'task_type': 'tool', 'component_identifier': 'my_tool'}`
>>> parse_task_key("auth:client-a:task456:tool:my_tool")
`{'task_scope': 'client-a', 'client_task_id': 'task456', 'task_type': 'tool', 'component_identifier': 'my_tool'}`
>>> parse_task_key("session123:task456:resource:file%3A%2F%2Fdata.txt")
`{'session_id': 'session123', 'client_task_id': 'task456', 'task_type': 'resource', 'component_identifier': 'file://data.txt'}`
>>> parse_task_key("anon:task456:tool:my_tool")
`{'task_scope': None, 'client_task_id': 'task456', 'task_type': 'tool', 'component_identifier': 'my_tool'}`
"""
parts = task_key.split(":", 3)
if len(parts) != 4:
tag, _, rest = task_key.partition(":")
if tag not in _VALID_TAGS or not rest:
raise ValueError(
f"Invalid task key format: {task_key}. "
f"Expected: {{session_id}}:{{client_task_id}}:{{task_type}}:{{component_identifier}}"
f"Expected leading tag in {_VALID_TAGS}."
)
if tag == _ANON_TAG:
parts = rest.split(":", 2)
if len(parts) != 3:
raise ValueError(
f"Invalid anonymous task key: {task_key}. "
f"Expected: anon:{{client_task_id}}:{{task_type}}:{{component_identifier}}"
)
client_task_id, task_type, encoded_identifier = parts
return {
"task_scope": None,
"client_task_id": client_task_id,
"task_type": task_type,
"component_identifier": unquote(encoded_identifier),
}
parts = rest.split(":", 3)
if len(parts) != 4:
raise ValueError(
f"Invalid authenticated task key: {task_key}. "
f"Expected: auth:{{enc_scope}}:{{client_task_id}}:{{task_type}}:{{component_identifier}}"
)
encoded_scope, client_task_id, task_type, encoded_identifier = parts
return {
"session_id": parts[0],
"client_task_id": parts[1],
"task_type": parts[2],
"component_identifier": unquote(parts[3]),
"task_scope": unquote(encoded_scope),
"client_task_id": client_task_id,
"task_type": task_type,
"component_identifier": unquote(encoded_identifier),
}
@ -82,10 +141,25 @@ def get_client_task_id_from_key(task_key: str) -> str:
task_key: Full encoded task key
Returns:
Client-provided task ID (second segment)
Client-provided task ID
Example:
>>> get_client_task_id_from_key("session123:task456:tool:my_tool")
Examples:
>>> get_client_task_id_from_key("auth:client-a:task456:tool:my_tool")
'task456'
>>> get_client_task_id_from_key("anon:task456:tool:my_tool")
'task456'
"""
return task_key.split(":", 3)[1]
return parse_task_key(task_key)["client_task_id"]
def task_redis_prefix(task_scope: str | None) -> str:
"""Return the Redis key prefix that owns a given scope.
Authenticated tasks live under ``fastmcp:task:auth:{enc_scope}``;
anonymous tasks live under ``fastmcp:task:anon``. Callers append
``f":{task_id}:..."`` to compose the final key.
"""
if task_scope is None:
return f"fastmcp:task:{_ANON_TAG}"
return f"fastmcp:task:{_AUTH_TAG}:{quote(task_scope, safe='')}"

10
src/fastmcp/server/tasks/notifications.py
View file

@ -211,10 +211,18 @@ async def _send_mcp_notification(
"input_required notification missing taskId, skipping relay"
)
return
if "task_scope" not in related_task:
logger.warning(
"input_required notification for task %s missing task_scope "
"metadata, skipping elicitation relay",
task_id,
)
return
task_scope = related_task["task_scope"]
from fastmcp.server.tasks.elicitation import relay_elicitation
task = asyncio.create_task(
relay_elicitation(session, session_id, task_id, elicitation, fastmcp),
relay_elicitation(session, task_scope, task_id, elicitation, fastmcp),
name=f"elicitation-relay-{task_id[:8]}",
)
_background_tasks.add(task)

42
src/fastmcp/server/tasks/requests.py
View file

@ -29,7 +29,8 @@ from fastmcp.prompts.base import Prompt
from fastmcp.resources.base import Resource
from fastmcp.resources.template import ResourceTemplate
from fastmcp.server.tasks.config import DEFAULT_POLL_INTERVAL_MS, DEFAULT_TTL_MS
from fastmcp.server.tasks.keys import parse_task_key
from fastmcp.server.tasks.context import get_task_scope
from fastmcp.server.tasks.keys import parse_task_key, task_redis_prefix
from fastmcp.tools.base import Tool
from fastmcp.utilities.versions import VersionSpec
@ -70,7 +71,7 @@ def _parse_key_version(key_suffix: str) -> tuple[str, str | None]:
async def _lookup_task_execution(
docket: Any,
session_id: str,
task_scope: str | None,
client_task_id: str,
) -> tuple[Any, str | None, int]:
"""Look up task execution and metadata from Redis.
@ -80,7 +81,7 @@ async def _lookup_task_execution(
Args:
docket: Docket instance
session_id: Session ID
task_scope: Authorization scope
client_task_id: Client-provided task ID
Returns:
@ -89,13 +90,10 @@ async def _lookup_task_execution(
Raises:
McpError: If task not found or execution not found
"""
task_meta_key = docket.key(f"fastmcp:task:{session_id}:{client_task_id}")
created_at_key = docket.key(
f"fastmcp:task:{session_id}:{client_task_id}:created_at"
)
poll_interval_key = docket.key(
f"fastmcp:task:{session_id}:{client_task_id}:poll_interval"
)
prefix = task_redis_prefix(task_scope)
task_meta_key = docket.key(f"{prefix}:{client_task_id}")
created_at_key = docket.key(f"{prefix}:{client_task_id}:created_at")
poll_interval_key = docket.key(f"{prefix}:{client_task_id}:poll_interval")
# Fetch metadata (single round-trip with mget)
async with docket.redis() as redis:
@ -144,7 +142,7 @@ async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskR
Returns:
GetTaskResult: Task status response with spec-compliant fields
"""
async with fastmcp.server.context.Context(fastmcp=server) as ctx:
async with fastmcp.server.context.Context(fastmcp=server):
client_task_id = params.get("taskId")
if not client_task_id:
raise McpError(
@ -153,8 +151,8 @@ async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskR
)
)
# Get session ID from Context
session_id = ctx.session_id
# Get authorization scope for task lookup
task_scope = get_task_scope()
# Get Docket instance
docket = server._docket
@ -168,7 +166,7 @@ async def tasks_get_handler(server: FastMCP, params: dict[str, Any]) -> GetTaskR
# Look up task execution and metadata
execution, created_at, poll_interval_ms = await _lookup_task_execution(
docket, session_id, client_task_id
docket, task_scope, client_task_id
)
# Sync state from Redis
@ -231,7 +229,7 @@ async def tasks_result_handler(server: FastMCP, params: dict[str, Any]) -> Any:
Returns:
MCP result (CallToolResult, GetPromptResult, or ReadResourceResult)
"""
async with fastmcp.server.context.Context(fastmcp=server) as ctx:
async with fastmcp.server.context.Context(fastmcp=server):
client_task_id = params.get("taskId")
if not client_task_id:
raise McpError(
@ -240,8 +238,8 @@ async def tasks_result_handler(server: FastMCP, params: dict[str, Any]) -> Any:
)
)
# Get session ID from Context
session_id = ctx.session_id
# Get authorization scope for task lookup
task_scope = get_task_scope()
# Get execution from Docket (use instance attribute for cross-task access)
docket = server._docket
@ -254,7 +252,7 @@ async def tasks_result_handler(server: FastMCP, params: dict[str, Any]) -> Any:
)
# Look up full task key from Redis
task_meta_key = docket.key(f"fastmcp:task:{session_id}:{client_task_id}")
task_meta_key = docket.key(f"{task_redis_prefix(task_scope)}:{client_task_id}")
async with docket.redis() as redis:
task_key_bytes = await redis.get(task_meta_key)
@ -432,7 +430,7 @@ async def tasks_cancel_handler(
Returns:
CancelTaskResult: Task status response showing cancelled state
"""
async with fastmcp.server.context.Context(fastmcp=server) as ctx:
async with fastmcp.server.context.Context(fastmcp=server):
client_task_id = params.get("taskId")
if not client_task_id:
raise McpError(
@ -441,8 +439,8 @@ async def tasks_cancel_handler(
)
)
# Get session ID from Context
session_id = ctx.session_id
# Get authorization scope for task lookup
task_scope = get_task_scope()
# Get Docket instance
docket = server._docket
@ -456,7 +454,7 @@ async def tasks_cancel_handler(
# Look up task execution and metadata
execution, created_at, poll_interval_ms = await _lookup_task_execution(
docket, session_id, client_task_id
docket, task_scope, client_task_id
)
# Cancel via Docket (now sets CANCELLED state natively)

14
src/fastmcp/server/tasks/subscriptions.py
View file

@ -16,7 +16,7 @@ from docket.execution import ExecutionState
from mcp.types import TaskStatusNotification, TaskStatusNotificationParams
from fastmcp.server.tasks.config import DEFAULT_TTL_MS
from fastmcp.server.tasks.keys import parse_task_key
from fastmcp.server.tasks.keys import parse_task_key, task_redis_prefix
from fastmcp.server.tasks.requests import DOCKET_TO_MCP_STATE
from fastmcp.utilities.logging import get_logger
@ -115,11 +115,11 @@ async def _send_status_notification(
state_map = DOCKET_TO_MCP_STATE
mcp_status = state_map.get(state, "failed")
# Extract session_id from task_key for Redis lookup
# Extract task_scope from task_key for Redis lookup
key_parts = parse_task_key(task_key)
session_id = key_parts["session_id"]
task_scope = key_parts["task_scope"]
created_at_key = docket.key(f"fastmcp:task:{session_id}:{task_id}:created_at")
created_at_key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:created_at")
async with docket.redis() as redis:
created_at_bytes = await redis.get(created_at_key)
@ -189,11 +189,11 @@ async def _send_progress_notification(
state_map = DOCKET_TO_MCP_STATE
mcp_status = state_map.get(execution.state, "failed")
# Extract session_id from task_key for Redis lookup
# Extract task_scope from task_key for Redis lookup
key_parts = parse_task_key(task_key)
session_id = key_parts["session_id"]
task_scope = key_parts["task_scope"]
created_at_key = docket.key(f"fastmcp:task:{session_id}:{task_id}:created_at")
created_at_key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:created_at")
async with docket.redis() as redis:
created_at_bytes = await redis.get(created_at_key)

33
tests/server/tasks/test_context_background_task.py
View file

@ -23,18 +23,22 @@ from fastmcp.client.elicitation import ElicitResult
from fastmcp.dependencies import CurrentDocket
from fastmcp.server.auth import AccessToken
from fastmcp.server.context import Context
from fastmcp.server.dependencies import (
TaskContextInfo,
TaskContextSnapshot,
_set_cached_snapshot,
get_access_token,
)
from fastmcp.server.dependencies import get_access_token
from fastmcp.server.elicitation import (
AcceptedElicitation,
CancelledElicitation,
DeclinedElicitation,
)
from fastmcp.server.tasks.context import (
TaskContextInfo,
TaskContextSnapshot,
_set_cached_snapshot,
get_task_scope,
)
from fastmcp.server.tasks.elicitation import handle_task_input
from fastmcp.server.tasks.keys import (
task_redis_prefix,
)
# =============================================================================
# Unit tests: Context API surface (no Redis/Docket needed)
@ -262,7 +266,8 @@ class TestBackgroundTaskIntegration:
assert origin != ""
# Verify the snapshot in Redis contains the same value
key = docket.key(f"fastmcp:task:{ctx.session_id}:{ctx.task_id}:snapshot")
task_scope = get_task_scope()
key = docket.key(f"{task_redis_prefix(task_scope)}:{ctx.task_id}:snapshot")
async with docket.redis() as redis:
raw = await redis.get(key)
@ -361,7 +366,7 @@ class TestBackgroundTaskIntegration:
# Task already completed — no elicitation waiting
success = await handle_task_input(
task_id=task.task_id,
session_id="nonexistent-session",
task_scope="nonexistent-scope",
action="accept",
content={"value": "too late"},
fastmcp=mcp,
@ -429,9 +434,9 @@ class TestAccessTokenInBackgroundTasks:
"test-task",
TaskContextSnapshot(access_token_json=expired.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
):
assert get_access_token() is None
@ -447,9 +452,9 @@ class TestAccessTokenInBackgroundTasks:
"test-task",
TaskContextSnapshot(access_token_json=valid.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
):
result = get_access_token()
assert result is not None
@ -466,9 +471,9 @@ class TestAccessTokenInBackgroundTasks:
"test-task",
TaskContextSnapshot(access_token_json=no_expiry.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
):
result = get_access_token()
assert result is not None

170
tests/server/tasks/test_task_keys.py Normal file
View file

@ -0,0 +1,170 @@
"""Tests for ``fastmcp.server.tasks.keys`` — the encoding boundary that
separates authenticated and anonymous task keyspaces.
Cross-scope isolation depends on these encodings being unambiguous and
round-trippable, so the tests cover: tag dispatch (``auth``/``anon``),
the ``None`` ⇄ anonymous round trip, encoding of values that contain the
``:`` delimiter, error paths for malformed keys, and the parity between
the Docket-key prefix and the Redis-key prefix.
"""
import pytest
from fastmcp.server.tasks.keys import (
build_task_key,
get_client_task_id_from_key,
parse_task_key,
task_redis_prefix,
)
ROUND_TRIP_CASES = [
("client-a", "task-1", "tool", "my_tool"),
(None, "task-1", "tool", "my_tool"),
("client-a", "task-1", "resource", "file://data.txt"),
(None, "task-1", "resource", "file://data.txt"),
("client-a", "task-1", "template", "users://{id}"),
("client-a", "task-1", "prompt", "greet@1.0.0"),
# Scope contains the inner separator used by get_task_scope (client_id|sub).
("client|sub-42", "task-1", "tool", "my_tool"),
# Adversarial: scope is literally the anon tag — must not collide.
("anon", "task-1", "tool", "my_tool"),
# Adversarial: scope is literally the legacy "_" sentinel.
("_", "task-1", "tool", "my_tool"),
# Scope contains every delimiter we care about.
("a:b/c d%e|f", "task-1", "tool", "my_tool"),
# Component identifier with colons, slashes, percent, spaces.
("client-a", "task-1", "resource", "https://x/y?z=1&q=a b"),
# UUID-shaped task id (the realistic case).
("client-a", "0c3e9b14-3a3f-4b3a-9b1a-1d8d6e6e0c11", "tool", "t"),
]
@pytest.mark.parametrize(
("scope", "task_id", "task_type", "identifier"), ROUND_TRIP_CASES
)
def test_round_trip_preserves_all_fields(
scope: str | None, task_id: str, task_type: str, identifier: str
):
key = build_task_key(scope, task_id, task_type, identifier)
parsed = parse_task_key(key)
assert parsed == {
"task_scope": scope,
"client_task_id": task_id,
"task_type": task_type,
"component_identifier": identifier,
}
@pytest.mark.parametrize(
("scope", "task_id", "task_type", "identifier"), ROUND_TRIP_CASES
)
def test_get_client_task_id_round_trip(
scope: str | None, task_id: str, task_type: str, identifier: str
):
key = build_task_key(scope, task_id, task_type, identifier)
assert get_client_task_id_from_key(key) == task_id
def test_authenticated_key_uses_auth_tag():
key = build_task_key("client-a", "task-1", "tool", "my_tool")
assert key.startswith("auth:")
assert key == "auth:client-a:task-1:tool:my_tool"
def test_anonymous_key_uses_anon_tag():
key = build_task_key(None, "task-1", "tool", "my_tool")
assert key.startswith("anon:")
assert key == "anon:task-1:tool:my_tool"
def test_anonymous_and_literal_anon_scope_have_disjoint_keyspaces():
"""A real anonymous task and a (hostile) authenticated task whose scope
literally equals "anon" must not collide."""
anon_key = build_task_key(None, "task-1", "tool", "x")
impostor_key = build_task_key("anon", "task-1", "tool", "x")
assert anon_key != impostor_key
assert parse_task_key(anon_key)["task_scope"] is None
assert parse_task_key(impostor_key)["task_scope"] == "anon"
def test_legacy_underscore_scope_is_just_a_string_now():
"""Belt-and-suspenders: a client_id of "_" no longer aliases anonymous."""
underscore_key = build_task_key("_", "task-1", "tool", "x")
anon_key = build_task_key(None, "task-1", "tool", "x")
assert underscore_key != anon_key
assert parse_task_key(underscore_key)["task_scope"] == "_"
def test_component_identifier_with_colons_is_recovered():
key = build_task_key("client-a", "task-1", "resource", "file://data:special.txt")
assert parse_task_key(key)["component_identifier"] == "file://data:special.txt"
def test_scope_with_colons_is_recovered():
key = build_task_key("a:b:c", "task-1", "tool", "t")
parsed = parse_task_key(key)
assert parsed["task_scope"] == "a:b:c"
assert parsed["client_task_id"] == "task-1"
def test_scope_pipe_separator_is_preserved():
"""``get_task_scope`` composes ``client_id|sub`` — the ``|`` must survive."""
key = build_task_key("client-a|user-42", "task-1", "tool", "t")
assert parse_task_key(key)["task_scope"] == "client-a|user-42"
@pytest.mark.parametrize(
"bad_key",
[
"",
"client-a:task-1:tool:my_tool", # legacy untagged format
"weird:client-a:task-1:tool:my_tool", # unknown tag
"auth:client-a:task-1:tool", # missing identifier
"auth:client-a", # truncated
"anon:task-1:tool", # truncated anon
"anon", # tag only
"auth", # tag only
":task-1:tool:t", # empty tag
],
)
def test_parse_rejects_malformed_keys(bad_key: str):
with pytest.raises(ValueError):
parse_task_key(bad_key)
def test_redis_prefix_authenticated():
assert task_redis_prefix("client-a") == "fastmcp:task:auth:client-a"
def test_redis_prefix_anonymous():
assert task_redis_prefix(None) == "fastmcp:task:anon"
def test_redis_prefix_disjoint_for_anon_vs_literal_anon_scope():
assert task_redis_prefix(None) != task_redis_prefix("anon")
def test_redis_prefix_disjoint_for_anon_vs_literal_underscore_scope():
assert task_redis_prefix(None) != task_redis_prefix("_")
def test_redis_prefix_encodes_special_characters():
# Colons, slashes, pipes in the scope must not break the prefix shape.
prefix = task_redis_prefix("client:a/b|sub")
assert prefix.startswith("fastmcp:task:auth:")
# Exactly four ":" delimiters: fastmcp / task / auth / encoded-scope.
assert prefix.count(":") == 3
def test_docket_and_redis_prefixes_agree_on_partition():
"""The Docket key tag and the Redis prefix tag must always match — that is
the load-bearing invariant for cross-scope isolation."""
auth_docket = build_task_key("client-a", "task-1", "tool", "x")
auth_redis = task_redis_prefix("client-a")
assert auth_docket.split(":", 1)[0] == "auth"
assert ":auth:" in auth_redis
anon_docket = build_task_key(None, "task-1", "tool", "x")
anon_redis = task_redis_prefix(None)
assert anon_docket.split(":", 1)[0] == "anon"
assert anon_redis.endswith(":anon")

144
tests/server/tasks/test_task_security.py
View file

@ -1,22 +1,27 @@
"""
Tests for session-based task ID isolation (CRITICAL SECURITY).
Tests for authorization-based task isolation (CRITICAL SECURITY).
Ensures that tasks are properly scoped to sessions and clients cannot
access each other's tasks.
Ensures that tasks are properly scoped to authorization identity and clients
cannot access each other's tasks.
"""
import pytest
from mcp.server.auth.middleware.auth_context import (
auth_context_var,
)
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.server.auth import AccessToken
@pytest.fixture
async def task_server():
def task_server():
"""Create a server with background tasks enabled."""
mcp = FastMCP("security-test-server")
@mcp.tool(task=True) # Enable background execution
@mcp.tool(task=True)
async def secret_tool(data: str) -> str:
"""A tool that processes sensitive data."""
return f"Secret result: {data}"
@ -24,24 +29,121 @@ async def task_server():
return mcp
async def test_same_session_can_access_all_its_tasks(task_server):
"""A single session can access all tasks it created."""
async def test_same_client_can_access_all_its_tasks(task_server: FastMCP):
"""A single authenticated client can access all tasks it created."""
token = AccessToken(
token="token-a",
client_id="client-a",
scopes=["read"],
)
reset = auth_context_var.set(AuthenticatedUser(token))
try:
async with Client(task_server) as client:
task1 = await client.call_tool(
"secret_tool", {"data": "first"}, task=True, task_id="task-1"
)
task2 = await client.call_tool(
"secret_tool", {"data": "second"}, task=True, task_id="task-2"
)
await task1.wait(timeout=2.0)
await task2.wait(timeout=2.0)
result1 = await task1.result()
result2 = await task2.result()
assert "first" in str(result1.data)
assert "second" in str(result2.data)
finally:
auth_context_var.reset(reset)
async def test_unauthenticated_client_can_access_its_tasks(task_server: FastMCP):
"""An unauthenticated client can access tasks it created (by task ID)."""
async with Client(task_server) as client:
# Submit multiple tasks
task1 = await client.call_tool(
"secret_tool", {"data": "first"}, task=True, task_id="task-1"
)
task2 = await client.call_tool(
"secret_tool", {"data": "second"}, task=True, task_id="task-2"
task = await client.call_tool(
"secret_tool", {"data": "hello"}, task=True, task_id="my-task"
)
await task.wait(timeout=2.0)
result = await task.result()
assert "hello" in str(result.data)
# Wait for both to complete
await task1.wait(timeout=2.0)
await task2.wait(timeout=2.0)
# Should be able to access both
result1 = await task1.result()
result2 = await task2.result()
def _set_auth(client_id: str, sub: str | None = None):
"""Install an auth context for a given client_id/sub. Returns the reset token."""
claims = {"sub": sub} if sub else {}
token = AccessToken(
token=f"token-{client_id}-{sub or ''}",
client_id=client_id,
scopes=["read"],
claims=claims,
)
return auth_context_var.set(AuthenticatedUser(token))
assert "first" in str(result1.data)
assert "second" in str(result2.data)
async def _submit_task_id(client: Client, data: str) -> str:
"""Submit a background task and return its server-assigned task id."""
task = await client.call_tool("secret_tool", {"data": data}, task=True)
await task.wait(timeout=2.0)
return task.task_id
async def test_distinct_clients_cannot_access_each_others_tasks(
task_server: FastMCP,
):
"""Two distinct authenticated clients live in disjoint scopes — looking up
a peer's task id returns 'not found'."""
reset = _set_auth("client-a")
try:
async with Client(task_server) as client_a:
task_id = await _submit_task_id(client_a, "client-a-secret")
finally:
auth_context_var.reset(reset)
reset = _set_auth("client-b")
try:
async with Client(task_server) as client_b:
with pytest.raises(Exception, match="not found"):
await client_b.get_task_status(task_id)
finally:
auth_context_var.reset(reset)
async def test_distinct_subs_same_client_id_cannot_access_each_others_tasks(
task_server: FastMCP,
):
"""Fixed-OAuth case: two users share a client_id but have distinct ``sub``
claims. The ``sub``-aware scope must still isolate them."""
shared_client = "shared-oauth-app"
reset = _set_auth(shared_client, sub="user-alice")
try:
async with Client(task_server) as alice:
task_id = await _submit_task_id(alice, "alice-secret")
finally:
auth_context_var.reset(reset)
reset = _set_auth(shared_client, sub="user-bob")
try:
async with Client(task_server) as bob:
with pytest.raises(Exception, match="not found"):
await bob.get_task_status(task_id)
finally:
auth_context_var.reset(reset)
async def test_authenticated_and_anonymous_keyspaces_are_disjoint(
task_server: FastMCP,
):
"""An anonymous client must not be able to read an authenticated client's
tasks (and vice versa) even when colliding on task id."""
reset = _set_auth("client-a")
try:
async with Client(task_server) as authed:
authed_task_id = await _submit_task_id(authed, "authed-secret")
finally:
auth_context_var.reset(reset)
async with Client(task_server) as anon:
with pytest.raises(Exception, match="not found"):
await anon.get_task_status(authed_task_id)