mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-20 04:24:17 +02:00
Scope tasks to authorization context, not session (#3800)
* 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
GitHub
parent
a6bd66aca8
commit
9a063963f4
11 changed files with 954 additions and 511 deletions
|
|
@ -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
389
src/fastmcp/server/tasks/context.py
Normal 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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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='')}"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
170
tests/server/tasks/test_task_keys.py
Normal 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")
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue