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

Restore task snapshot via a worker-level dependency (#3945)

Browse source 
* Restore task snapshot via a worker-level dependency

`get_access_token()` returned `None` inside background tasks whenever
`FASTMCP_DOCKET_URL` pointed at a `redis+cluster://` URL. The write side
was fine — it went through `docket.redis()`, which is cluster-aware —
but fastmcp kept a parallel sync Redis client just to read the snapshot
back, and `Redis.from_url()` rejects the cluster scheme.

Docket 0.19.1 ships worker-level dependencies that resolve per task in
the same asyncio.Task as user code, so ContextVars propagate cleanly.
That lets us load the snapshot once via `restore_task_snapshot` and
drop the sync Redis path entirely. Sync helpers like
`get_access_token()` and `get_http_request()` now just read a
ContextVar; Docket is the sole Redis consumer.

Closes #3897

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Revert TaskKey stub to a plain return

NotImplementedError would fire at module import if anything evaluated
the default; a no-op stub keeps the module usable without the
fastmcp[tasks] extra, which is what we want.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Chris Guidry 2026-04-17 08:24:42 -04:00 • committed by GitHub
commit 34313ea112
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 207 additions and 147 deletions
Download patch file Download diff file Expand all files Collapse all files

3
pyproject.toml
View file

@ -60,7 +60,7 @@ azure = ["azure-identity>=1.16.0", "PyJWT>=2.12.0"]
code-mode = ["pydantic-monty==0.0.11"]
gemini = ["google-genai>=1.18.0"]
openai = ["openai>=1.102.0"]
tasks = ["pydocket>=0.19.0"]
tasks = ["pydocket>=0.19.1"]
[dependency-groups]
dev = [
@ -118,7 +118,6 @@ style = "pep440"
bump = true
fallback-version = "0.0.0"
[tool.pytest.ini_options]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"

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

@ -81,13 +81,10 @@ __all__ = [
# 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,
_recall_snapshot,
get_task_context,
get_task_server,
get_task_session,
@ -378,10 +375,12 @@ def get_http_request() -> Request:
if request is None:
request = _current_http_request.get()
# In Docket workers, restore a minimal request from the snapshotted headers.
# Uses sync fallback chain: ContextVar → in-memory dict → sync Redis.
# In Docket workers, restore a minimal request from the snapshotted
# headers. The snapshot is preloaded by restore_task_snapshot before
# user code runs, so this is a pure ContextVar read.
if request is None:
snapshot = _get_task_snapshot_sync()
task_info = get_task_context()
snapshot = _recall_snapshot(task_info.task_id) if task_info else None
task_headers = snapshot.http_headers if snapshot else None
if task_headers:
request = Request(
@ -495,11 +494,12 @@ def get_access_token() -> AccessToken | None:
if access_token is None:
access_token = _sdk_get_access_token()
# Fall back to background task snapshot (#3095)
# In Docket workers, neither HTTP request nor SDK context var are available.
# Uses sync fallback chain: ContextVar → in-memory dict → sync Redis.
# Fall back to background task snapshot (#3095). In Docket workers,
# neither the HTTP request nor the SDK context var is available; the
# snapshot is preloaded by restore_task_snapshot before user code runs.
if access_token is None:
snapshot = _get_task_snapshot_sync()
task_info = get_task_context()
snapshot = _recall_snapshot(task_info.task_id) if task_info else None
if snapshot is not None and snapshot.access_token_json is not None:
task_token = AccessToken.model_validate_json(snapshot.access_token_json)
if task_token.expires_at is not None:
@ -757,10 +757,10 @@ class _CurrentContext(Dependency["Context"]):
if task_info is not None:
server = get_server()
# Load unified snapshot (sets _task_snapshot ContextVar)
snapshot = await _load_task_snapshot_async(
task_info.task_scope, task_info.task_id
)
# The snapshot is preloaded by restore_task_snapshot (worker-level
# Docket dependency) before any task code runs, so this is a pure
# ContextVar read — no Redis I/O here.
snapshot = _recall_snapshot(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

14
src/fastmcp/server/mixins/lifespan.py
View file

@ -77,13 +77,14 @@ class LifespanMixin:
return
# Docket is available AND there are task-enabled components
from docket import Docket, Worker
from docket import Depends, Docket, Worker
from fastmcp import settings
from fastmcp.server.dependencies import (
_current_docket,
_current_worker,
)
from fastmcp.server.tasks.context import restore_task_snapshot
# Create Docket instance using configured name and URL
async with Docket(
@ -108,8 +109,15 @@ class LifespanMixin:
if settings.docket.worker_name:
worker_kwargs["name"] = settings.docket.worker_name
# Create and start Worker
async with Worker(docket, **worker_kwargs) as worker:
# Create and start Worker. The restore_task_snapshot
# worker-level dependency runs before every task so the
# per-task snapshot ContextVar is populated before user
# code or task-scoped dependencies observe it.
async with Worker(
docket,
dependencies=[Depends(restore_task_snapshot)],
**worker_kwargs,
) as worker:
self._worker = worker
worker_token = _current_worker.set(worker)
try:

181
src/fastmcp/server/tasks/context.py
View file

@ -14,10 +14,22 @@ import weakref
from collections import OrderedDict
from contextvars import ContextVar
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING
from fastmcp.server.tasks.keys import parse_task_key, task_redis_prefix
try:
from docket import TaskKey
except ImportError:
def TaskKey() -> str: # type: ignore[no-redef]
# Stub so this module stays importable without the fastmcp[tasks]
# extra. ``restore_task_snapshot`` is only ever invoked inside a
# Docket worker, where the real ``docket.TaskKey`` sentinel is
# always present.
return ""
if TYPE_CHECKING:
from docket import Docket
from mcp.server.session import ServerSession
@ -170,7 +182,7 @@ class TaskContextSnapshot:
ttl_seconds: int,
) -> None:
"""Store this snapshot as a single Redis key."""
key = docket.key(_snapshot_redis_key(task_scope, task_id))
key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot")
async with docket.redis() as redis:
await redis.set(key, self.to_json(), ex=ttl_seconds)
@ -182,13 +194,21 @@ _task_snapshot: ContextVar[tuple[str, TaskContextSnapshot] | None] = ContextVar(
)
def _set_cached_snapshot(task_id: str, snapshot: TaskContextSnapshot) -> None:
"""Cache a snapshot keyed by task_id."""
def _remember_snapshot(task_id: str, snapshot: TaskContextSnapshot) -> None:
"""Bind a snapshot to the current asyncio context under ``task_id``.
Nothing outside this task's context sees it; stale entries left in a
reused context are ignored on recall.
"""
_task_snapshot.set((task_id, snapshot))
def _get_cached_snapshot(task_id: str) -> TaskContextSnapshot | None:
"""Get cached snapshot if it belongs to this task."""
def _recall_snapshot(task_id: str) -> TaskContextSnapshot | None:
"""Return the snapshot bound for ``task_id`` in the current context.
Returns ``None`` if nothing is bound, or if the bound entry belongs to
a different task (a stale leftover from a reused asyncio context).
"""
cached = _task_snapshot.get()
if cached is not None:
cached_task_id, snapshot = cached
@ -197,21 +217,35 @@ def _get_cached_snapshot(task_id: str) -> TaskContextSnapshot | None:
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"
def get_task_session_id() -> str | None:
"""Get the session_id for the current background task, if available.
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.
Reads the cached snapshot set by the worker-level restore dependency.
Returns None if not in a task context or the snapshot wasn't restored.
"""
cached = _get_cached_snapshot(task_id)
if cached is not None:
return cached
task_info = get_task_context()
if task_info is None:
return None
snapshot = _recall_snapshot(task_info.task_id)
return snapshot.session_id if snapshot else None
async def restore_task_snapshot(key: str = TaskKey()) -> None:
"""Worker-level Docket dependency that restores the task-context snapshot.
Runs before each fastmcp-owned task, populating the snapshot ContextVar
so user code — and any task-scoped dependency like ``_CurrentContext`` —
sees a ready snapshot without touching Redis itself. All Redis I/O
goes through Docket's async client, so cluster URLs and the memory://
backend work transparently (#3897). Failures are non-fatal: the task
still runs, and sync helpers return ``None`` as they would have before
the snapshot was captured.
"""
try:
parts = parse_task_key(key)
except ValueError:
# Non-fastmcp key (e.g. docket scheduler internals) — nothing to do.
return
from fastmcp.server.dependencies import _current_docket, get_server
@ -222,109 +256,20 @@ async def _load_task_snapshot_async(
if docket is None:
docket = _current_docket.get()
if docket is None:
return None
return
task_scope = parts["task_scope"]
task_id = parts["client_task_id"]
try:
async with docket.redis() as redis:
raw = await redis.get(docket.key(_snapshot_redis_key(task_scope, task_id)))
raw = await redis.get(
docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot")
)
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)
return
_remember_snapshot(task_id, TaskContextSnapshot.from_json(raw))
except Exception:
_logger.warning("Failed to restore task snapshot for %s", key, exc_info=True)
# In-process optimization: when the Docket worker runs in the same process as

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

@ -32,7 +32,7 @@ from fastmcp.server.elicitation import (
from fastmcp.server.tasks.context import (
TaskContextInfo,
TaskContextSnapshot,
_set_cached_snapshot,
_remember_snapshot,
get_task_scope,
)
from fastmcp.server.tasks.elicitation import handle_task_input
@ -430,13 +430,13 @@ class TestAccessTokenInBackgroundTasks:
scopes=["read"],
expires_at=int(datetime.now(timezone.utc).timestamp()) - 3600,
)
_set_cached_snapshot(
_remember_snapshot(
"test-task",
TaskContextSnapshot(access_token_json=expired.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
):
assert get_access_token() is None
@ -448,13 +448,13 @@ class TestAccessTokenInBackgroundTasks:
scopes=["read"],
expires_at=int(datetime.now(timezone.utc).timestamp()) + 3600,
)
_set_cached_snapshot(
_remember_snapshot(
"test-task",
TaskContextSnapshot(access_token_json=valid.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
):
result = get_access_token()
assert result is not None
@ -467,13 +467,13 @@ class TestAccessTokenInBackgroundTasks:
client_id="test-client",
scopes=["read"],
)
_set_cached_snapshot(
_remember_snapshot(
"test-task",
TaskContextSnapshot(access_token_json=no_expiry.model_dump_json()),
)
fake_ctx = TaskContextInfo(task_id="test-task", task_scope="s")
with patch(
"fastmcp.server.tasks.context.get_task_context", return_value=fake_ctx
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
):
result = get_access_token()
assert result is not None

108
tests/server/tasks/test_snapshot_restore.py Normal file
View file

@ -0,0 +1,108 @@
"""Tests for ``restore_task_snapshot`` — the worker-level Docket dependency
that restores the task-context snapshot into the ``_task_snapshot``
ContextVar before each task runs.
With the snapshot restored up front, sync helpers (``get_access_token``,
``get_http_request``, etc.) never need to hit Redis themselves. These
tests exercise the restore path end-to-end (via in-memory Docket) and
the edge cases around non-fastmcp keys and failed restores.
"""
from __future__ import annotations
from unittest.mock import patch
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
from fastmcp.server.dependencies import get_access_token
from fastmcp.server.tasks.context import (
TaskContextSnapshot,
_recall_snapshot,
get_task_context,
restore_task_snapshot,
)
async def test_snapshot_restored_before_user_code_runs():
"""A tool with no declared deps finds the snapshot already cached."""
mcp = FastMCP("snapshot-restore-test")
seen_cached: list[bool] = []
@mcp.tool(task=True)
async def bare_tool() -> str:
info = get_task_context()
assert info is not None
seen_cached.append(_recall_snapshot(info.task_id) is not None)
return "ok"
async with Client(mcp) as client:
task = await client.call_tool("bare_tool", {}, task=True)
await task.result()
assert seen_cached == [True]
async def test_get_access_token_in_bg_task_without_context_dep():
"""Issue #3897 repro: get_access_token() works in a bg task that does
not declare Context as a dependency."""
mcp = FastMCP("access-token-test")
seen_tokens: list[str | None] = []
@mcp.tool(task=True)
async def bare_tool() -> str:
token = get_access_token()
seen_tokens.append(token.token if token else None)
return "ok"
test_token = AccessToken(
token="jwt-3897",
client_id="test-client",
scopes=["read"],
claims={"sub": "user-x"},
)
auth_context_var.set(AuthenticatedUser(test_token))
async with Client(mcp) as client:
task = await client.call_tool("bare_tool", {}, task=True)
await task.result()
assert seen_tokens == ["jwt-3897"]
async def test_restore_failure_is_nonfatal():
"""If deserialization blows up, the task still runs to completion and
the snapshot cache stays empty."""
mcp = FastMCP("restore-failure-test")
seen_cached: list[bool] = []
@mcp.tool(task=True)
async def bare_tool() -> str:
info = get_task_context()
assert info is not None
seen_cached.append(_recall_snapshot(info.task_id) is not None)
return "ok"
def boom(*_args, **_kwargs):
raise RuntimeError("simulated deserialization failure")
async with Client(mcp) as client:
with patch.object(TaskContextSnapshot, "from_json", boom):
task = await client.call_tool("bare_tool", {}, task=True)
result = await task.result()
assert result.data == "ok"
assert seen_cached == [False]
async def test_restore_skipped_for_non_fastmcp_task_keys():
"""The restore dep returns cleanly for keys it doesn't recognize and
writes nothing to the snapshot cache."""
# Direct calls bypass the worker, so Redis/Docket never gets involved
# — any attempt to touch them would raise.
await restore_task_snapshot(key="not-a-fastmcp-key")
await restore_task_snapshot(key="weird:client-a:task-1:tool:my_tool")
await restore_task_snapshot(key="")

8
uv.lock generated
View file

@ -929,7 +929,7 @@ requires-dist = [
{ name = "py-key-value-aio", extras = ["filetree", "keyring", "memory"], specifier = ">=0.4.4,<0.5.0" },
{ name = "pydantic", extras = ["email"], specifier = ">=2.11.7" },
{ name = "pydantic-monty", marker = "extra == 'code-mode'", specifier = "==0.0.11" },
{ name = "pydocket", marker = "extra == 'tasks'", specifier = ">=0.19.0" },
{ name = "pydocket", marker = "extra == 'tasks'", specifier = ">=0.19.1" },
{ name = "pyjwt", marker = "extra == 'azure'", specifier = ">=2.12.0" },
{ name = "pyperclip", specifier = ">=1.9.0" },
{ name = "python-dotenv", specifier = ">=1.1.0" },
@ -2287,7 +2287,7 @@ wheels = [
[[package]]
name = "pydocket"
version = "0.19.0"
version = "0.19.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "cloudpickle" },
@ -2306,9 +2306,9 @@ dependencies = [
{ name = "tzdata", marker = "sys_platform == 'win32'" },
{ name = "uncalled-for" },
]
sdist = { url = "https://files.pythonhosted.org/packages/98/6e/0db603ce4d82072b1a61798340e408ec04b3a77647f537881ff5b93c31f6/pydocket-0.19.0.tar.gz", hash = "sha256:00bff620d80cd2fad34ccbbe526dce24a9de8cdc1d2b94d305739668a98e308a", size = 355531, upload-time = "2026-04-10T17:25:38.112Z" }
sdist = { url = "https://files.pythonhosted.org/packages/dd/6e/cdfef4cf65c568deea932fc0c9b521b0287c62cb36a25f66739da5f7528d/pydocket-0.19.1.tar.gz", hash = "sha256:83135a2c171f6600c6ab6d4e0739bf76496b87383ab9a2172e3bdfcd52bc8b05", size = 363300, upload-time = "2026-04-15T21:10:47.16Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/55/46/7bed93ecff9015c4a8dcabfaab3d490b45ec8e5847b30ac9671b9c01def8/pydocket-0.19.0-py3-none-any.whl", hash = "sha256:8531e64b989673a17d055ee4498ca8c3505310c5af4e7fd09c7b00fb2f29aa19", size = 99271, upload-time = "2026-04-10T17:25:36.657Z" },
{ url = "https://files.pythonhosted.org/packages/7a/e1/488d5e35b69e1d0c222dea76f3ed63f7237aff8b995d1d7a2d869b929bd0/pydocket-0.19.1-py3-none-any.whl", hash = "sha256:ec17f73452856482959ab90265ac8630347f2cbd21f0edb782016ae57faad1d1", size = 100981, upload-time = "2026-04-15T21:10:45.551Z" },
]
[[package]]