From 34313ea1126aba0cfb48f1439f5e39e4f3c57bc4 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Fri, 17 Apr 2026 08:24:42 -0400 Subject: [PATCH] Restore task snapshot via a worker-level dependency (#3945) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 * 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 --------- Co-authored-by: Claude Opus 4.6 --- pyproject.toml | 3 +- src/fastmcp/server/dependencies.py | 30 +-- src/fastmcp/server/mixins/lifespan.py | 14 +- src/fastmcp/server/tasks/context.py | 181 ++++++------------ .../tasks/test_context_background_task.py | 14 +- tests/server/tasks/test_snapshot_restore.py | 108 +++++++++++ uv.lock | 8 +- 7 files changed, 209 insertions(+), 149 deletions(-) create mode 100644 tests/server/tasks/test_snapshot_restore.py diff --git a/pyproject.toml b/pyproject.toml index d3748bbb4..d06445a1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index eb1f24c3b..5035e0918 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -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 diff --git a/src/fastmcp/server/mixins/lifespan.py b/src/fastmcp/server/mixins/lifespan.py index 508041009..b64d725ed 100644 --- a/src/fastmcp/server/mixins/lifespan.py +++ b/src/fastmcp/server/mixins/lifespan.py @@ -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: diff --git a/src/fastmcp/server/tasks/context.py b/src/fastmcp/server/tasks/context.py index 667d2b213..7ade43dfc 100644 --- a/src/fastmcp/server/tasks/context.py +++ b/src/fastmcp/server/tasks/context.py @@ -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 diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py index 4b7f66583..40e9ddb90 100644 --- a/tests/server/tasks/test_context_background_task.py +++ b/tests/server/tasks/test_context_background_task.py @@ -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 diff --git a/tests/server/tasks/test_snapshot_restore.py b/tests/server/tasks/test_snapshot_restore.py new file mode 100644 index 000000000..946e1bf1b --- /dev/null +++ b/tests/server/tasks/test_snapshot_restore.py @@ -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="") diff --git a/uv.lock b/uv.lock index 79922ca9c..ce8a12af0 100644 --- a/uv.lock +++ b/uv.lock @@ -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]]