diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index 447b8e7c7..eb1f24c3b 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -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, diff --git a/src/fastmcp/server/tasks/context.py b/src/fastmcp/server/tasks/context.py new file mode 100644 index 000000000..667d2b213 --- /dev/null +++ b/src/fastmcp/server/tasks/context.py @@ -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 diff --git a/src/fastmcp/server/tasks/elicitation.py b/src/fastmcp/server/tasks/elicitation.py index cc6ac2624..7b5d76dfd 100644 --- a/src/fastmcp/server/tasks/elicitation.py +++ b/src/fastmcp/server/tasks/elicitation.py @@ -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, diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index 2fe9eb83f..aa9d73bdc 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/src/fastmcp/server/tasks/handlers.py @@ -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) diff --git a/src/fastmcp/server/tasks/keys.py b/src/fastmcp/server/tasks/keys.py index 0e28cf592..10af6a6f9 100644 --- a/src/fastmcp/server/tasks/keys.py +++ b/src/fastmcp/server/tasks/keys.py @@ -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='')}" diff --git a/src/fastmcp/server/tasks/notifications.py b/src/fastmcp/server/tasks/notifications.py index 6656bc361..852b2bb37 100644 --- a/src/fastmcp/server/tasks/notifications.py +++ b/src/fastmcp/server/tasks/notifications.py @@ -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) diff --git a/src/fastmcp/server/tasks/requests.py b/src/fastmcp/server/tasks/requests.py index 8743356e5..6c062de0c 100644 --- a/src/fastmcp/server/tasks/requests.py +++ b/src/fastmcp/server/tasks/requests.py @@ -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) diff --git a/src/fastmcp/server/tasks/subscriptions.py b/src/fastmcp/server/tasks/subscriptions.py index 772b82671..37a4bf2ea 100644 --- a/src/fastmcp/server/tasks/subscriptions.py +++ b/src/fastmcp/server/tasks/subscriptions.py @@ -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) diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py index ad472410f..4b7f66583 100644 --- a/tests/server/tasks/test_context_background_task.py +++ b/tests/server/tasks/test_context_background_task.py @@ -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 diff --git a/tests/server/tasks/test_task_keys.py b/tests/server/tasks/test_task_keys.py new file mode 100644 index 000000000..06a64f8f1 --- /dev/null +++ b/tests/server/tasks/test_task_keys.py @@ -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") diff --git a/tests/server/tasks/test_task_security.py b/tests/server/tasks/test_task_security.py index 88af3aa7d..5d3b16ffa 100644 --- a/tests/server/tasks/test_task_security.py +++ b/tests/server/tasks/test_task_security.py @@ -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)