mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
Unify background task context forwarding, fix concurrent dependency bugs (#3710)
* Unify background task context forwarding and fix concurrent dependency bugs We've been getting a steady trickle of edge-case reports around background tasks and contextual dependencies over the last few months (#3654, #3656, #3569). Each one pointed at a different symptom, but they all traced back to the same area: the way context is negotiated between the "frontend" server and Docket workers was grown piecemeal, with each new piece of context (access tokens, HTTP headers, origin request IDs) getting its own Redis key, its own restore function, and its own ContextVar. This made it hard to reason about what state was available where, and the shared-instance Dependency pattern made concurrent tasks stomp on each other's cleanup state. This takes a step back and reworks the whole thing as a single unified system: - Dependency subclasses (_CurrentContext, Progress, _CurrentAccessToken, etc.) are now stateless factories — __aenter__ returns a fresh per-invocation object, so concurrent tasks never share mutable state. Fixes #3654, #3656. - The three individual context-snapshot Redis keys (access_token, http_headers, origin_request_id) are collapsed into a single TaskContextSnapshot stored as one JSON key per task. The three _restore_task_* functions and two ContextVars they populated are gone. - Sync functions like get_http_request() and get_access_token() now find the snapshot transparently in background tasks via a 3-tier sync fallback: ContextVar (set by _CurrentContext for functions with deps) → in-memory dict (same-process workers) → sync Redis GET (out-of-process workers). No function wrapping needed. - The _wrap_for_task_http_headers hack is deleted. FunctionTool registers its raw function with Docket so Docket sees and resolves ALL dependencies, including Docket-native ones like Retry and Timeout. - ProxyTool.from_mcp_tool() now propagates execution.taskSupport metadata from remote tools. Fixes #3569. - Removed redundant _current_docket/_current_worker ContextVar management from Context.__aenter__/__aexit__ (they're only set in the lifespan now). Closes #3654 Closes #3656 Closes #3569 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Address code review feedback - _OptionalCurrentContext: guard __aexit__ against cleaning up contexts it didn't create (check is_background_task before delegating) - Narrow except clauses in snapshot loading (OSError, JSONDecodeError, etc. instead of bare Exception) - Fix docstrings on register_with_docket for resources/prompts/templates - Simplify Progress: read ExecutionProgress directly from current_execution instead of creating and manually entering a DocketProgress wrapper 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Use pop-on-access transfer buffer instead of bounded LRU cache for snapshots The in-memory snapshot dict is a transfer mechanism, not a cache. Entries go in at submission and come out at the worker's first access. Using pop instead of get means the dict only holds entries during the brief submission-to-execution window, bounded by task concurrency (~10) rather than a 10,000-entry LRU limit. 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Drop in-memory transfer buffer, use sync Redis for all backends Instead of maintaining an in-memory dict to bridge the async/sync gap, use a sync Redis client directly. For memory:// backends (fakeredis), shares the same FakeServer instance via docket._redis.get_memory_server() so data written by the async Docket client is visible to sync reads. For real Redis, creates a standard sync connection. No in-process state to manage at all. 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Move snapshot operations to TaskContextSnapshot methods capture(), from_json(), to_json(), save() are now classmethod/instance methods on the dataclass instead of free functions. Deduplicates JSON parsing that was copy-pasted between the async and sync load paths. 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Trim implementation details from register_with_docket docstrings 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Clarify docket lookup comment in submit_to_docket 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Restore docket/worker ContextVar bridge in Context.__aenter__ Servers that own the Docket (the parent) re-set _current_docket/_current_worker from their instance attributes when entering a Context. Mounted children skip this (their _docket is None), so they inherit the parent's value. This is needed for ASGI deployments where ContextVars set during the lifespan don't propagate to request handlers. 🤖 Generated with Claude Code Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Key snapshot cache by task_id to prevent cross-task context leakage Docket workers may reuse the same asyncio context for sequential tasks. The ContextVar cache now stores (task_id, snapshot) tuples so stale entries from previous tasks are automatically ignored. 🤖 Generated with Claude Code 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
697426b39e
commit
d41bcb2c9e
13 changed files with 586 additions and 349 deletions
|
|
@ -333,11 +333,7 @@ class FunctionPrompt(Prompt):
|
|||
raise PromptError(f"Error rendering prompt {self.name}.") from e
|
||||
|
||||
def register_with_docket(self, docket: Docket) -> None:
|
||||
"""Register this prompt with docket for background execution.
|
||||
|
||||
FunctionPrompt registers the underlying function, which has the user's
|
||||
Depends parameters for docket to resolve.
|
||||
"""
|
||||
"""Register this prompt with docket for background execution."""
|
||||
if not self.task_config.supports_tasks():
|
||||
return
|
||||
docket.register(self.fn, names=[self.key])
|
||||
|
|
|
|||
|
|
@ -228,11 +228,7 @@ class FunctionResource(Resource):
|
|||
return result
|
||||
|
||||
def register_with_docket(self, docket: Docket) -> None:
|
||||
"""Register this resource with docket for background execution.
|
||||
|
||||
FunctionResource registers the underlying function, which has the user's
|
||||
Depends parameters for docket to resolve.
|
||||
"""
|
||||
"""Register this resource with docket for background execution."""
|
||||
if not self.task_config.supports_tasks():
|
||||
return
|
||||
docket.register(self.fn, names=[self.key])
|
||||
|
|
|
|||
|
|
@ -439,11 +439,7 @@ class FunctionResourceTemplate(ResourceTemplate):
|
|||
return result
|
||||
|
||||
def register_with_docket(self, docket: Docket) -> None:
|
||||
"""Register this template with docket for background execution.
|
||||
|
||||
FunctionResourceTemplate registers the underlying function, which has the
|
||||
user's Depends parameters for docket to resolve.
|
||||
"""
|
||||
"""Register this template with docket for background execution."""
|
||||
if not self.task_config.supports_tasks():
|
||||
return
|
||||
docket.register(self.fn, names=[self.key])
|
||||
|
|
|
|||
|
|
@ -272,16 +272,18 @@ class Context:
|
|||
|
||||
self._server_token = _current_server.set(weakref.ref(self.fastmcp))
|
||||
|
||||
# Set docket/worker from server instance for this request's context.
|
||||
# This ensures ContextVars work even in ASGI environments (Lambda, FastAPI mount)
|
||||
# where lifespan ContextVars don't propagate to request handlers.
|
||||
server = self.fastmcp
|
||||
# Re-set docket/worker from the server instance so mounted children
|
||||
# inherit the parent's Docket via the ContextVar. Only servers that
|
||||
# own the Docket (the parent) have _docket set; children skip this,
|
||||
# leaving the parent's value in place.
|
||||
if is_docket_available():
|
||||
server = self.fastmcp
|
||||
if server._docket is not None:
|
||||
self._docket_token = _current_docket.set(server._docket)
|
||||
if server._worker is not None:
|
||||
self._worker_token = _current_worker.set(server._worker)
|
||||
else:
|
||||
|
||||
if not is_docket_available():
|
||||
# Without docket, the lifespan won't provide a SharedContext,
|
||||
# so create one scoped to this Context for Shared() dependencies.
|
||||
self._shared_context = SharedContext()
|
||||
|
|
@ -297,7 +299,6 @@ class Context:
|
|||
_current_worker,
|
||||
)
|
||||
|
||||
# Mirror __aenter__: clean up docket/worker tokens or SharedContext
|
||||
if hasattr(self, "_worker_token"):
|
||||
_current_worker.reset(self._worker_token)
|
||||
del self._worker_token
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import weakref
|
|||
from collections import OrderedDict
|
||||
from collections.abc import AsyncGenerator, Callable
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from contextvars import ContextVar, Token
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from functools import lru_cache
|
||||
|
|
@ -65,6 +65,7 @@ __all__ = [
|
|||
"CurrentWorker",
|
||||
"Progress",
|
||||
"TaskContextInfo",
|
||||
"TaskContextSnapshot",
|
||||
"TokenClaim",
|
||||
"get_access_token",
|
||||
"get_context",
|
||||
|
|
@ -205,14 +206,215 @@ def register_task_server(task_id: str, server: FastMCP) -> None:
|
|||
|
||||
_current_docket: ContextVar[Docket | None] = ContextVar("docket", default=None)
|
||||
_current_worker: ContextVar[Worker | None] = ContextVar("worker", default=None)
|
||||
_task_access_token: ContextVar[AccessToken | None] = ContextVar(
|
||||
"task_access_token", default=None
|
||||
)
|
||||
_task_http_headers: ContextVar[dict[str, str] | None] = ContextVar(
|
||||
"task_http_headers", 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
|
||||
|
|
@ -459,8 +661,10 @@ def get_http_request() -> Request:
|
|||
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.
|
||||
if request is None:
|
||||
task_headers = _task_http_headers.get()
|
||||
snapshot = _get_task_snapshot_sync()
|
||||
task_headers = snapshot.http_headers if snapshot else None
|
||||
if task_headers:
|
||||
request = Request(
|
||||
{
|
||||
|
|
@ -575,12 +779,11 @@ def get_access_token() -> AccessToken | None:
|
|||
|
||||
# Fall back to background task snapshot (#3095)
|
||||
# In Docket workers, neither HTTP request nor SDK context var are available.
|
||||
# The token was snapshotted in Redis at submit_to_docket() time and restored
|
||||
# into this ContextVar by _CurrentContext.__aenter__().
|
||||
# Uses sync fallback chain: ContextVar → in-memory dict → sync Redis.
|
||||
if access_token is None:
|
||||
task_token = _task_access_token.get()
|
||||
if task_token is not None:
|
||||
# Check expiration: if expires_at is set and past, treat as expired
|
||||
snapshot = _get_task_snapshot_sync()
|
||||
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:
|
||||
if task_token.expires_at < int(datetime.now(timezone.utc).timestamp()):
|
||||
return None
|
||||
|
|
@ -811,112 +1014,17 @@ async def resolve_dependencies(
|
|||
# so that get_dependency_parameters can detect them.
|
||||
|
||||
|
||||
async def _restore_task_access_token(
|
||||
session_id: str, task_id: str
|
||||
) -> Token[AccessToken | None] | None:
|
||||
"""Restore the access token snapshot from Redis into a ContextVar.
|
||||
|
||||
Called when setting up context in a Docket worker. The token was stored at
|
||||
submit_to_docket() time. The token is restored regardless of expiration;
|
||||
get_access_token() checks expiry when reading from the ContextVar.
|
||||
|
||||
Returns:
|
||||
The ContextVar token for resetting, or None if nothing was restored.
|
||||
"""
|
||||
docket = _current_docket.get()
|
||||
if docket is None:
|
||||
return None
|
||||
|
||||
token_key = docket.key(f"fastmcp:task:{session_id}:{task_id}:access_token")
|
||||
try:
|
||||
async with docket.redis() as redis:
|
||||
token_data = await redis.get(token_key)
|
||||
if token_data is not None:
|
||||
restored = AccessToken.model_validate_json(token_data)
|
||||
return _task_access_token.set(restored)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"Failed to restore access token for task %s:%s",
|
||||
session_id,
|
||||
task_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def _restore_task_http_headers(
|
||||
session_id: str, task_id: str
|
||||
) -> Token[dict[str, str] | None] | None:
|
||||
"""Restore the HTTP header snapshot from Redis into a ContextVar."""
|
||||
docket = _current_docket.get()
|
||||
if docket is None:
|
||||
return None
|
||||
|
||||
headers_key = docket.key(f"fastmcp:task:{session_id}:{task_id}:http_headers")
|
||||
try:
|
||||
async with docket.redis() as redis:
|
||||
headers_data = await redis.get(headers_key)
|
||||
if headers_data is None:
|
||||
return None
|
||||
if isinstance(headers_data, bytes):
|
||||
headers_data = headers_data.decode()
|
||||
restored = json.loads(str(headers_data))
|
||||
if not isinstance(restored, dict):
|
||||
return None
|
||||
return _task_http_headers.set(
|
||||
{str(name).lower(): str(value) for name, value in restored.items()}
|
||||
)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"Failed to restore HTTP headers for task %s:%s",
|
||||
session_id,
|
||||
task_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def _restore_task_origin_request_id(session_id: str, task_id: str) -> str | None:
|
||||
"""Restore the origin request ID snapshot for a background task.
|
||||
|
||||
Returns None if no request ID was captured at submission time.
|
||||
"""
|
||||
docket = _current_docket.get()
|
||||
if docket is None:
|
||||
return None
|
||||
|
||||
request_id_key = docket.key(
|
||||
f"fastmcp:task:{session_id}:{task_id}:origin_request_id"
|
||||
)
|
||||
try:
|
||||
async with docket.redis() as redis:
|
||||
request_id_data = await redis.get(request_id_key)
|
||||
if request_id_data is None:
|
||||
return None
|
||||
if isinstance(request_id_data, bytes):
|
||||
return request_id_data.decode()
|
||||
return str(request_id_data)
|
||||
except Exception:
|
||||
_logger.warning(
|
||||
"Failed to restore origin request ID for task %s:%s",
|
||||
session_id,
|
||||
task_id,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class _CurrentContext(Dependency["Context"]):
|
||||
"""Async context manager for Context dependency.
|
||||
|
||||
In foreground (request) mode: returns the active context from _current_context.
|
||||
In background (Docket worker) mode: creates a task-aware Context with task_id
|
||||
and restores the access token snapshot from Redis.
|
||||
"""
|
||||
and loads the unified task snapshot from Redis.
|
||||
|
||||
_context: Context | None = None
|
||||
_access_token_cv_token: Token[AccessToken | None] | None = None
|
||||
_http_headers_cv_token: Token[dict[str, str] | None] | None = None
|
||||
The shared default instance is a stateless factory. All per-invocation
|
||||
state lives on the returned Context or in task-local ContextVars, so
|
||||
concurrent tasks never share mutable state.
|
||||
"""
|
||||
|
||||
async def __aenter__(self) -> Context:
|
||||
from fastmcp.server.context import Context, _current_context
|
||||
|
|
@ -929,36 +1037,24 @@ class _CurrentContext(Dependency["Context"]):
|
|||
# Check if we're in a Docket worker context
|
||||
task_info = get_task_context()
|
||||
if task_info is not None:
|
||||
# Get session from registry (registered when task was submitted)
|
||||
session = get_task_session(task_info.session_id)
|
||||
# Get server from ContextVar
|
||||
server = get_server()
|
||||
origin_request_id = await _restore_task_origin_request_id(
|
||||
|
||||
# Load unified snapshot (sets _task_snapshot ContextVar)
|
||||
snapshot = await _load_task_snapshot_async(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
# Create task-aware Context
|
||||
self._context = Context(
|
||||
origin_request_id = snapshot.origin_request_id if snapshot else None
|
||||
|
||||
ctx = Context(
|
||||
fastmcp=server,
|
||||
session=session,
|
||||
task_id=task_info.task_id,
|
||||
origin_request_id=origin_request_id,
|
||||
)
|
||||
# Enter the context to set up ContextVars
|
||||
await self._context.__aenter__()
|
||||
await ctx.__aenter__()
|
||||
return ctx
|
||||
|
||||
# Restore access token snapshot from Redis (#3095)
|
||||
self._access_token_cv_token = await _restore_task_access_token(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
|
||||
# Restore HTTP headers snapshot from Redis (#3631)
|
||||
self._http_headers_cv_token = await _restore_task_http_headers(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
|
||||
return self._context
|
||||
|
||||
# Neither foreground nor background context available
|
||||
raise RuntimeError(
|
||||
"No active context found. This can happen if:\n"
|
||||
" - Called outside an MCP request handler\n"
|
||||
|
|
@ -972,40 +1068,29 @@ class _CurrentContext(Dependency["Context"]):
|
|||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
# Clean up access token ContextVar
|
||||
if self._access_token_cv_token is not None:
|
||||
_task_access_token.reset(self._access_token_cv_token)
|
||||
self._access_token_cv_token = None
|
||||
# Clean up HTTP headers ContextVar
|
||||
if self._http_headers_cv_token is not None:
|
||||
_task_http_headers.reset(self._http_headers_cv_token)
|
||||
self._http_headers_cv_token = None
|
||||
# Clean up if we created a context for background task
|
||||
if self._context is not None:
|
||||
await self._context.__aexit__(exc_type, exc_value, traceback)
|
||||
self._context = None
|
||||
from fastmcp.server.context import _current_context
|
||||
|
||||
ctx = _current_context.get()
|
||||
if ctx is not None and ctx.is_background_task:
|
||||
await ctx.__aexit__(exc_type, exc_value, traceback)
|
||||
|
||||
|
||||
class _OptionalCurrentContext(Dependency["Context | None"]):
|
||||
"""Context dependency that degrades to None when no context is active.
|
||||
"""Context dependency that returns None instead of raising when no context
|
||||
is active. Used for ``ctx: Context = None`` parameter patterns.
|
||||
|
||||
This is implemented as a wrapper (composition), not a subclass of
|
||||
`_CurrentContext`, to avoid overriding `__aenter__` with an incompatible
|
||||
return type.
|
||||
Delegates entirely to ``_CurrentContext`` — just catches the RuntimeError.
|
||||
Cleanup is handled by ``_CurrentContext.__aexit__`` reading from the
|
||||
task-local ContextVar.
|
||||
"""
|
||||
|
||||
_inner: _CurrentContext | None = None
|
||||
|
||||
async def __aenter__(self) -> Context | None:
|
||||
inner = _CurrentContext()
|
||||
try:
|
||||
context = await inner.__aenter__()
|
||||
return await _CurrentContext().__aenter__()
|
||||
except RuntimeError as exc:
|
||||
if "No active context found" in str(exc):
|
||||
return None
|
||||
raise
|
||||
self._inner = inner
|
||||
return context
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
|
|
@ -1013,10 +1098,11 @@ class _OptionalCurrentContext(Dependency["Context | None"]):
|
|||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
if self._inner is None:
|
||||
return
|
||||
await self._inner.__aexit__(exc_type, exc_value, traceback)
|
||||
self._inner = None
|
||||
from fastmcp.server.context import _current_context
|
||||
|
||||
ctx = _current_context.get()
|
||||
if ctx is not None and ctx.is_background_task:
|
||||
await _CurrentContext().__aexit__(exc_type, exc_value, traceback)
|
||||
|
||||
|
||||
def CurrentContext() -> Context:
|
||||
|
|
@ -1054,7 +1140,14 @@ class _CurrentDocket(Dependency["Docket"]):
|
|||
|
||||
async def __aenter__(self) -> Docket:
|
||||
require_docket("CurrentDocket()")
|
||||
docket = _current_docket.get()
|
||||
# Check server instance first, fall back to ContextVar for mounted children
|
||||
# whose parent owns the Docket
|
||||
try:
|
||||
docket = get_server()._docket
|
||||
except RuntimeError:
|
||||
docket = None
|
||||
if docket is None:
|
||||
docket = _current_docket.get()
|
||||
if docket is None:
|
||||
raise RuntimeError(
|
||||
"No Docket instance found. Docket is only initialized when there are "
|
||||
|
|
@ -1104,7 +1197,13 @@ class _CurrentWorker(Dependency["Worker"]):
|
|||
|
||||
async def __aenter__(self) -> Worker:
|
||||
require_docket("CurrentWorker()")
|
||||
worker = _current_worker.get()
|
||||
# Check server instance first, fall back to ContextVar for mounted children
|
||||
try:
|
||||
worker = get_server()._worker
|
||||
except RuntimeError:
|
||||
worker = None
|
||||
if worker is None:
|
||||
worker = _current_worker.get()
|
||||
if worker is None:
|
||||
raise RuntimeError(
|
||||
"No Worker instance found. Worker is only initialized when there are "
|
||||
|
|
@ -1191,20 +1290,8 @@ def CurrentFastMCP() -> FastMCP:
|
|||
class _CurrentRequest(Dependency[Request]):
|
||||
"""Async context manager for HTTP Request dependency."""
|
||||
|
||||
_task_http_headers_cv_token: Token[dict[str, str] | None] | None = None
|
||||
|
||||
async def __aenter__(self) -> Request:
|
||||
try:
|
||||
return get_http_request()
|
||||
except RuntimeError:
|
||||
task_info = get_task_context()
|
||||
if task_info is None:
|
||||
raise
|
||||
if _task_http_headers.get() is None:
|
||||
self._task_http_headers_cv_token = await _restore_task_http_headers(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
return get_http_request()
|
||||
return get_http_request()
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
|
|
@ -1212,9 +1299,7 @@ class _CurrentRequest(Dependency[Request]):
|
|||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
if self._task_http_headers_cv_token is not None:
|
||||
_task_http_headers.reset(self._task_http_headers_cv_token)
|
||||
self._task_http_headers_cv_token = None
|
||||
pass
|
||||
|
||||
|
||||
def CurrentRequest() -> Request:
|
||||
|
|
@ -1246,15 +1331,7 @@ def CurrentRequest() -> Request:
|
|||
class _CurrentHeaders(Dependency[dict[str, str]]):
|
||||
"""Async context manager for HTTP Headers dependency."""
|
||||
|
||||
_task_http_headers_cv_token: Token[dict[str, str] | None] | None = None
|
||||
|
||||
async def __aenter__(self) -> dict[str, str]:
|
||||
if _task_http_headers.get() is None:
|
||||
task_info = get_task_context()
|
||||
if task_info is not None:
|
||||
self._task_http_headers_cv_token = await _restore_task_http_headers(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
return get_http_headers(include={"authorization"})
|
||||
|
||||
async def __aexit__(
|
||||
|
|
@ -1263,9 +1340,7 @@ class _CurrentHeaders(Dependency[dict[str, str]]):
|
|||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
if self._task_http_headers_cv_token is not None:
|
||||
_task_http_headers.reset(self._task_http_headers_cv_token)
|
||||
self._task_http_headers_cv_token = None
|
||||
pass
|
||||
|
||||
|
||||
def CurrentHeaders() -> dict[str, str]:
|
||||
|
|
@ -1388,16 +1463,14 @@ class InMemoryProgress:
|
|||
|
||||
|
||||
class Progress(Dependency["Progress"]):
|
||||
"""FastMCP Progress dependency that works in both server and worker contexts.
|
||||
"""Progress dependency that works in both server and worker contexts.
|
||||
|
||||
Handles three execution modes:
|
||||
- In Docket worker: Uses the execution's progress (observable via Redis)
|
||||
- In FastMCP server with Docket: Falls back to in-memory progress
|
||||
- In FastMCP server without Docket: Uses in-memory progress
|
||||
In a Docket worker, delegates to the execution's Redis-backed progress
|
||||
(observable across processes). Otherwise, uses in-memory tracking.
|
||||
|
||||
This allows tools to use Progress() regardless of whether they're called
|
||||
immediately or as background tasks, and regardless of whether pydocket
|
||||
is installed.
|
||||
The shared default instance acts as a stateless factory — ``__aenter__``
|
||||
creates a fresh ``Progress`` per invocation so concurrent tasks never
|
||||
share mutable state.
|
||||
"""
|
||||
|
||||
_impl: ProgressLike | None = None
|
||||
|
|
@ -1407,18 +1480,19 @@ class Progress(Dependency["Progress"]):
|
|||
if server_ref is None or server_ref() is None:
|
||||
raise RuntimeError("Progress dependency requires a FastMCP server context.")
|
||||
|
||||
if is_docket_available():
|
||||
from docket.dependencies import Progress as DocketProgress
|
||||
instance = Progress()
|
||||
|
||||
if is_docket_available():
|
||||
try:
|
||||
docket_progress = DocketProgress()
|
||||
self._impl = await docket_progress.__aenter__()
|
||||
return self
|
||||
from docket.dependencies import current_execution
|
||||
|
||||
instance._impl = current_execution.get().progress
|
||||
return instance
|
||||
except LookupError:
|
||||
pass
|
||||
|
||||
self._impl = InMemoryProgress()
|
||||
return self
|
||||
instance._impl = InMemoryProgress()
|
||||
return instance
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
|
|
@ -1426,7 +1500,7 @@ class Progress(Dependency["Progress"]):
|
|||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
self._impl = None
|
||||
pass
|
||||
|
||||
@property
|
||||
def current(self) -> int | None:
|
||||
|
|
@ -1468,22 +1542,9 @@ class Progress(Dependency["Progress"]):
|
|||
class _CurrentAccessToken(Dependency[AccessToken]):
|
||||
"""Async context manager for AccessToken dependency."""
|
||||
|
||||
_access_token_cv_token: Token[AccessToken | None] | None = None
|
||||
|
||||
async def __aenter__(self) -> AccessToken:
|
||||
token = get_access_token()
|
||||
|
||||
# If no token found and we're in a Docket worker, try restoring from
|
||||
# Redis. This handles the case where ctx: Context is not in the
|
||||
# function signature, so _CurrentContext never ran the restoration.
|
||||
if token is None:
|
||||
task_info = get_task_context()
|
||||
if task_info is not None:
|
||||
self._access_token_cv_token = await _restore_task_access_token(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
token = get_access_token()
|
||||
|
||||
if token is None:
|
||||
raise RuntimeError(
|
||||
"No access token found. Ensure authentication is configured "
|
||||
|
|
@ -1497,9 +1558,7 @@ class _CurrentAccessToken(Dependency[AccessToken]):
|
|||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
if self._access_token_cv_token is not None:
|
||||
_task_access_token.reset(self._access_token_cv_token)
|
||||
self._access_token_cv_token = None
|
||||
pass
|
||||
|
||||
|
||||
def CurrentAccessToken() -> AccessToken:
|
||||
|
|
|
|||
|
|
@ -90,14 +90,12 @@ class LifespanMixin:
|
|||
name=settings.docket.name,
|
||||
url=settings.docket.url,
|
||||
) as docket:
|
||||
# Store on server instance for cross-task access (FastMCPTransport)
|
||||
self._docket = docket
|
||||
|
||||
# Register task-enabled components with Docket
|
||||
for component in task_components:
|
||||
component.register_with_docket(docket)
|
||||
|
||||
# Set Docket in ContextVar so CurrentDocket can access it
|
||||
docket_token = _current_docket.set(docket)
|
||||
try:
|
||||
# Build worker kwargs from settings
|
||||
|
|
@ -112,9 +110,7 @@ class LifespanMixin:
|
|||
|
||||
# Create and start Worker
|
||||
async with Worker(docket, **worker_kwargs) as worker:
|
||||
# Store on server instance for cross-context access
|
||||
self._worker = worker
|
||||
# Set Worker in ContextVar so CurrentWorker can access it
|
||||
worker_token = _current_worker.set(worker)
|
||||
try:
|
||||
worker_task = asyncio.create_task(worker.run_forever())
|
||||
|
|
@ -128,9 +124,7 @@ class LifespanMixin:
|
|||
_current_worker.reset(worker_token)
|
||||
self._worker = None
|
||||
finally:
|
||||
# Reset ContextVar
|
||||
_current_docket.reset(docket_token)
|
||||
# Clear instance attribute
|
||||
self._docket = None
|
||||
finally:
|
||||
# Reset server ContextVar
|
||||
|
|
|
|||
|
|
@ -108,6 +108,7 @@ class ProxyTool(Tool):
|
|||
icons=mcp_tool.icons,
|
||||
meta=mcp_tool.meta,
|
||||
tags=get_fastmcp_metadata(mcp_tool.meta).get("tags", []),
|
||||
execution=mcp_tool.execution,
|
||||
)
|
||||
|
||||
async def run(
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ Handles queuing tool/prompt/resource executions to Docket as background tasks.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -16,10 +15,9 @@ from mcp.shared.exceptions import McpError
|
|||
from mcp.types import INTERNAL_ERROR, ErrorData
|
||||
|
||||
from fastmcp.server.dependencies import (
|
||||
TaskContextSnapshot,
|
||||
_current_docket,
|
||||
get_access_token,
|
||||
get_context,
|
||||
get_http_headers,
|
||||
register_task_server,
|
||||
)
|
||||
from fastmcp.server.tasks.config import TaskMeta
|
||||
|
|
@ -78,7 +76,9 @@ async def submit_to_docket(
|
|||
except RuntimeError:
|
||||
session_id = "internal"
|
||||
|
||||
docket = _current_docket.get()
|
||||
# Try the server's own Docket first; fall back to the ContextVar for
|
||||
# mounted children (whose parent server owns the Docket instance).
|
||||
docket = ctx.fastmcp._docket or _current_docket.get()
|
||||
if docket is None:
|
||||
raise McpError(
|
||||
ErrorData(
|
||||
|
|
@ -111,36 +111,17 @@ async def submit_to_docket(
|
|||
poll_interval_key = docket.key(
|
||||
f"fastmcp:task:{session_id}:{server_task_id}:poll_interval"
|
||||
)
|
||||
origin_request_id_key = docket.key(
|
||||
f"fastmcp:task:{session_id}:{server_task_id}:origin_request_id"
|
||||
)
|
||||
poll_interval_ms = int(component.task_config.poll_interval.total_seconds() * 1000)
|
||||
origin_request_id = (
|
||||
str(ctx.request_context.request_id) if ctx.request_context is not None else None
|
||||
)
|
||||
|
||||
# Snapshot the current access token (if any) for background task access (#3095)
|
||||
access_token = get_access_token()
|
||||
access_token_key = docket.key(
|
||||
f"fastmcp:task:{session_id}:{server_task_id}:access_token"
|
||||
)
|
||||
http_headers = get_http_headers(include_all=True)
|
||||
http_headers_key = docket.key(
|
||||
f"fastmcp:task:{session_id}:{server_task_id}:http_headers"
|
||||
)
|
||||
# Snapshot all context (access token, headers, origin request ID) as a single key
|
||||
snapshot = TaskContextSnapshot.capture()
|
||||
|
||||
async with docket.redis() as redis:
|
||||
await redis.set(task_meta_key, task_key, ex=ttl_seconds)
|
||||
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)
|
||||
if origin_request_id is not None:
|
||||
await redis.set(origin_request_id_key, origin_request_id, ex=ttl_seconds)
|
||||
if access_token is not None:
|
||||
await redis.set(
|
||||
access_token_key, access_token.model_dump_json(), ex=ttl_seconds
|
||||
)
|
||||
if http_headers:
|
||||
await redis.set(http_headers_key, json.dumps(http_headers), ex=ttl_seconds)
|
||||
|
||||
await snapshot.save(docket, session_id, server_task_id, ttl_seconds)
|
||||
|
||||
# Register session for Context access in background workers (SEP-1686)
|
||||
# This enables elicitation/sampling from background tasks via weakref
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import inspect
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
|
|
@ -28,12 +27,7 @@ import fastmcp
|
|||
from fastmcp.decorators import resolve_task_config
|
||||
from fastmcp.exceptions import FastMCPDeprecationWarning
|
||||
from fastmcp.server.auth.authorization import AuthCheck
|
||||
from fastmcp.server.dependencies import (
|
||||
_restore_task_http_headers,
|
||||
_task_http_headers,
|
||||
get_task_context,
|
||||
without_injected_parameters,
|
||||
)
|
||||
from fastmcp.server.dependencies import without_injected_parameters
|
||||
from fastmcp.server.tasks.config import TaskConfig
|
||||
from fastmcp.tools.base import (
|
||||
Tool,
|
||||
|
|
@ -289,14 +283,13 @@ class FunctionTool(Tool):
|
|||
def register_with_docket(self, docket: Docket) -> None:
|
||||
"""Register this tool with docket for background execution.
|
||||
|
||||
FunctionTool registers the underlying function, which has the user's
|
||||
Depends parameters for docket to resolve. The function is wrapped to
|
||||
eagerly restore HTTP headers from Redis so that get_http_request()
|
||||
works even without explicit dependency injection.
|
||||
Registers the raw function so Docket sees and resolves ALL
|
||||
dependencies — both FastMCP's (CurrentContext, Progress) and
|
||||
Docket-native ones (Retry, Timeout, ConcurrencyLimit).
|
||||
"""
|
||||
if not self.task_config.supports_tasks():
|
||||
return
|
||||
docket.register(_wrap_for_task_http_headers(self.fn), names=[self.key])
|
||||
docket.register(self.fn, names=[self.key])
|
||||
|
||||
async def add_to_docket(
|
||||
self,
|
||||
|
|
@ -324,34 +317,6 @@ class FunctionTool(Tool):
|
|||
return await docket.add(lookup_key, **kwargs)(**arguments)
|
||||
|
||||
|
||||
def _wrap_for_task_http_headers(fn: Callable[..., Any]) -> Callable[..., Any]:
|
||||
"""Wrap a function to restore HTTP headers in background task workers.
|
||||
|
||||
Uses functools.wraps so docket sees the original signature for dependency
|
||||
resolution while the wrapper eagerly populates _task_http_headers before
|
||||
the user's function runs.
|
||||
"""
|
||||
|
||||
@functools.wraps(fn)
|
||||
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
||||
task_info = get_task_context()
|
||||
token = None
|
||||
if task_info is not None and _task_http_headers.get() is None:
|
||||
token = await _restore_task_http_headers(
|
||||
task_info.session_id, task_info.task_id
|
||||
)
|
||||
try:
|
||||
result = fn(*args, **kwargs)
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
return result
|
||||
finally:
|
||||
if token is not None:
|
||||
_task_http_headers.reset(token)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
@overload
|
||||
def tool(fn: F) -> F: ...
|
||||
@overload
|
||||
|
|
|
|||
213
tests/server/tasks/test_concurrent_dependencies.py
Normal file
213
tests/server/tasks/test_concurrent_dependencies.py
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
"""Tests for concurrent dependency resolution in foreground and background tasks.
|
||||
|
||||
Regression tests for:
|
||||
- #3654: ValueError when concurrent Docket tasks share a Dependency instance
|
||||
that stores a ContextVar token on `self`
|
||||
- #3656: Progress raises AssertionError when concurrent tasks share `_impl`
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.dependencies import Progress
|
||||
from fastmcp.server.context import Context
|
||||
from fastmcp.server.dependencies import (
|
||||
get_access_token,
|
||||
get_http_headers,
|
||||
)
|
||||
|
||||
|
||||
async def test_concurrent_foreground_tools_with_context():
|
||||
"""Multiple concurrent tool calls sharing the same CurrentContext() default
|
||||
should not raise ValueError from ContextVar token resets (#3654)."""
|
||||
mcp = FastMCP("test")
|
||||
results: list[str] = []
|
||||
|
||||
@mcp.tool()
|
||||
async def slow_tool(name: str, ctx: Context) -> str:
|
||||
await asyncio.sleep(0.05)
|
||||
results.append(name)
|
||||
return f"done:{name}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tasks = [client.call_tool("slow_tool", {"name": f"task-{i}"}) for i in range(4)]
|
||||
outcomes = await asyncio.gather(*tasks)
|
||||
|
||||
assert len(outcomes) == 4
|
||||
for outcome in outcomes:
|
||||
assert outcome.content[0].text.startswith("done:")
|
||||
|
||||
|
||||
async def test_concurrent_foreground_tools_with_progress():
|
||||
"""Multiple concurrent tool calls sharing the same Progress() default
|
||||
should not raise AssertionError from _impl being None (#3656)."""
|
||||
mcp = FastMCP("test")
|
||||
|
||||
@mcp.tool()
|
||||
async def variable_tool(
|
||||
name: str, delay: float, progress: Progress = Progress()
|
||||
) -> str:
|
||||
await progress.set_total(3)
|
||||
await progress.increment()
|
||||
await asyncio.sleep(delay)
|
||||
await progress.increment()
|
||||
await progress.set_message(f"finishing {name}")
|
||||
await progress.increment()
|
||||
return f"done:{name}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tasks = [
|
||||
client.call_tool(
|
||||
"variable_tool", {"name": f"t-{i}", "delay": 0.01 * (i + 1)}
|
||||
)
|
||||
for i in range(4)
|
||||
]
|
||||
outcomes = await asyncio.gather(*tasks)
|
||||
|
||||
assert len(outcomes) == 4
|
||||
for outcome in outcomes:
|
||||
assert outcome.content[0].text.startswith("done:")
|
||||
|
||||
|
||||
async def test_concurrent_background_tasks_with_context():
|
||||
"""Multiple concurrent background tasks sharing _CurrentContext() should
|
||||
not raise ValueError from ContextVar token resets (#3654)."""
|
||||
mcp = FastMCP("test")
|
||||
|
||||
@mcp.tool(task=True)
|
||||
async def bg_tool(name: str, ctx: Context) -> str:
|
||||
await asyncio.sleep(0.05)
|
||||
return f"bg:{name}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
task_handles = [
|
||||
await client.call_tool("bg_tool", {"name": f"bg-{i}"}, task=True)
|
||||
for i in range(4)
|
||||
]
|
||||
results = await asyncio.gather(*[t.result() for t in task_handles])
|
||||
|
||||
assert len(results) == 4
|
||||
for result in results:
|
||||
assert result.content[0].text.startswith("bg:")
|
||||
|
||||
|
||||
async def test_concurrent_background_tasks_with_progress():
|
||||
"""Multiple concurrent background tasks sharing Progress() should
|
||||
not raise AssertionError from _impl being None (#3656)."""
|
||||
mcp = FastMCP("test")
|
||||
|
||||
@mcp.tool(task=True)
|
||||
async def bg_progress_tool(
|
||||
name: str, delay: float, progress: Progress = Progress()
|
||||
) -> str:
|
||||
await progress.set_total(3)
|
||||
await progress.increment()
|
||||
await asyncio.sleep(delay)
|
||||
await progress.increment()
|
||||
await progress.set_message(f"bg finishing {name}")
|
||||
await progress.increment()
|
||||
return f"bg:{name}"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
task_handles = [
|
||||
await client.call_tool(
|
||||
"bg_progress_tool",
|
||||
{"name": f"bg-{i}", "delay": 0.01 * (i + 1)},
|
||||
task=True,
|
||||
)
|
||||
for i in range(4)
|
||||
]
|
||||
results = await asyncio.gather(*[t.result() for t in task_handles])
|
||||
|
||||
assert len(results) == 4
|
||||
for result in results:
|
||||
assert result.content[0].text.startswith("bg:")
|
||||
|
||||
|
||||
async def test_dependency_aenter_returns_fresh_instances():
|
||||
"""Verify that Dependency.__aenter__ returns independent per-invocation
|
||||
objects, not the shared default."""
|
||||
mcp = FastMCP("test")
|
||||
|
||||
instances: list[Context] = []
|
||||
|
||||
@mcp.tool()
|
||||
async def capture_context(ctx: Context) -> str:
|
||||
instances.append(ctx)
|
||||
return "ok"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
await asyncio.gather(
|
||||
client.call_tool("capture_context", {}),
|
||||
client.call_tool("capture_context", {}),
|
||||
)
|
||||
|
||||
assert len(instances) == 2
|
||||
assert instances[0] is not instances[1]
|
||||
|
||||
|
||||
async def test_progress_aenter_returns_fresh_instances():
|
||||
"""Verify that Progress.__aenter__ returns independent per-invocation
|
||||
objects, not the shared default."""
|
||||
progress_instances: list[Progress] = []
|
||||
|
||||
mcp = FastMCP("test")
|
||||
|
||||
@mcp.tool()
|
||||
async def capture_progress(progress: Progress = Progress()) -> str:
|
||||
progress_instances.append(progress)
|
||||
await progress.set_total(1)
|
||||
await progress.increment()
|
||||
return "ok"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
await asyncio.gather(
|
||||
client.call_tool("capture_progress", {}),
|
||||
client.call_tool("capture_progress", {}),
|
||||
)
|
||||
|
||||
assert len(progress_instances) == 2
|
||||
assert progress_instances[0] is not progress_instances[1]
|
||||
assert progress_instances[0]._impl is not progress_instances[1]._impl
|
||||
|
||||
|
||||
async def test_sync_context_functions_work_in_background_without_deps():
|
||||
"""Sync functions like get_http_request() should work in background tasks
|
||||
even when the tool declares no Context or CurrentRequest dependency.
|
||||
|
||||
This exercises the sync Redis fallback path (_get_task_snapshot_sync →
|
||||
_load_snapshot_sync_redis) which must work with both memory:// (fakeredis)
|
||||
and real Redis backends.
|
||||
"""
|
||||
mcp = FastMCP("test")
|
||||
|
||||
@mcp.tool(task=True)
|
||||
async def bare_sync_access() -> dict[str, str]:
|
||||
headers = get_http_headers()
|
||||
return {"has_headers": str(bool(headers))}
|
||||
|
||||
async with Client(mcp) as client:
|
||||
task = await client.call_tool("bare_sync_access", {}, task=True)
|
||||
result = await task.result()
|
||||
assert result.data == {"has_headers": "False"}
|
||||
|
||||
|
||||
async def test_sync_context_functions_work_in_background_with_context():
|
||||
"""Sync functions work via ContextVar when _CurrentContext loads the snapshot."""
|
||||
mcp = FastMCP("test")
|
||||
|
||||
@mcp.tool(task=True)
|
||||
async def context_sync_access(ctx: Context) -> dict[str, str]:
|
||||
headers = get_http_headers()
|
||||
token = get_access_token()
|
||||
return {
|
||||
"has_headers": str(bool(headers)),
|
||||
"has_token": str(token is not None),
|
||||
"is_background": str(ctx.is_background_task),
|
||||
}
|
||||
|
||||
async with Client(mcp) as client:
|
||||
task = await client.call_tool("context_sync_access", {}, task=True)
|
||||
result = await task.result()
|
||||
assert result.data["is_background"] == "True"
|
||||
|
|
@ -6,10 +6,16 @@ no mocking of Redis, Docket, or session internals.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from mcp import ServerSession
|
||||
from mcp.server.auth.middleware.auth_context import auth_context_var
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
||||
from pydantic import BaseModel
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
|
|
@ -17,8 +23,17 @@ 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 get_access_token
|
||||
from fastmcp.server.elicitation import AcceptedElicitation, DeclinedElicitation
|
||||
from fastmcp.server.dependencies import (
|
||||
TaskContextInfo,
|
||||
TaskContextSnapshot,
|
||||
_set_cached_snapshot,
|
||||
get_access_token,
|
||||
)
|
||||
from fastmcp.server.elicitation import (
|
||||
AcceptedElicitation,
|
||||
CancelledElicitation,
|
||||
DeclinedElicitation,
|
||||
)
|
||||
from fastmcp.server.tasks.elicitation import handle_task_input
|
||||
|
||||
# =============================================================================
|
||||
|
|
@ -119,10 +134,6 @@ class TestElicitFailFast:
|
|||
This test patches ONLY push_notification — all other components
|
||||
(Docket, Redis, session) are real via the memory:// backend.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastmcp.server.elicitation import CancelledElicitation
|
||||
|
||||
mcp = FastMCP("failfast-test")
|
||||
elicit_started = asyncio.Event()
|
||||
captured: dict[str, object] = {}
|
||||
|
|
@ -250,16 +261,16 @@ class TestBackgroundTaskIntegration:
|
|||
assert isinstance(origin, str)
|
||||
assert origin != ""
|
||||
|
||||
key = docket.key(
|
||||
f"fastmcp:task:{ctx.session_id}:{ctx.task_id}:origin_request_id"
|
||||
)
|
||||
# Verify the snapshot in Redis contains the same value
|
||||
key = docket.key(f"fastmcp:task:{ctx.session_id}:{ctx.task_id}:snapshot")
|
||||
async with docket.redis() as redis:
|
||||
raw = await redis.get(key)
|
||||
|
||||
assert raw is not None
|
||||
if isinstance(raw, bytes):
|
||||
raw = raw.decode()
|
||||
assert str(raw) == origin
|
||||
snapshot = json.loads(raw)
|
||||
assert snapshot["origin_request_id"] == origin
|
||||
return "ok"
|
||||
|
||||
async with Client(mcp) as client:
|
||||
|
|
@ -311,7 +322,6 @@ class TestBackgroundTaskIntegration:
|
|||
|
||||
async def test_elicit_with_pydantic_model(self):
|
||||
"""E2E: tool elicits structured Pydantic input via elicitation_handler."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class UserInfo(BaseModel):
|
||||
name: str
|
||||
|
|
@ -371,9 +381,6 @@ class TestAccessTokenInBackgroundTasks:
|
|||
|
||||
async def test_token_round_trips_through_background_task(self):
|
||||
"""E2E: token set at submit time is available inside the worker."""
|
||||
from mcp.server.auth.middleware.auth_context import auth_context_var
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
||||
|
||||
mcp = FastMCP("token-roundtrip")
|
||||
|
||||
@mcp.tool(task=True)
|
||||
|
|
@ -412,49 +419,60 @@ class TestAccessTokenInBackgroundTasks:
|
|||
|
||||
async def test_expired_token_returns_none(self):
|
||||
"""get_access_token() returns None when task token has expired."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastmcp.server.dependencies import _task_access_token
|
||||
|
||||
expired = AccessToken(
|
||||
token="expired-jwt",
|
||||
client_id="test-client",
|
||||
scopes=["read"],
|
||||
expires_at=int(datetime.now(timezone.utc).timestamp()) - 3600,
|
||||
)
|
||||
_task_access_token.set(expired)
|
||||
assert get_access_token() is None
|
||||
_set_cached_snapshot(
|
||||
"test-task",
|
||||
TaskContextSnapshot(access_token_json=expired.model_dump_json()),
|
||||
)
|
||||
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
|
||||
with patch(
|
||||
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
|
||||
):
|
||||
assert get_access_token() is None
|
||||
|
||||
async def test_valid_token_with_future_expiry(self):
|
||||
"""get_access_token() returns token when expiry is in the future."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastmcp.server.dependencies import _task_access_token
|
||||
|
||||
valid = AccessToken(
|
||||
token="valid-jwt",
|
||||
client_id="test-client",
|
||||
scopes=["read"],
|
||||
expires_at=int(datetime.now(timezone.utc).timestamp()) + 3600,
|
||||
)
|
||||
_task_access_token.set(valid)
|
||||
result = get_access_token()
|
||||
assert result is not None
|
||||
assert result.token == "valid-jwt"
|
||||
_set_cached_snapshot(
|
||||
"test-task",
|
||||
TaskContextSnapshot(access_token_json=valid.model_dump_json()),
|
||||
)
|
||||
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
|
||||
with patch(
|
||||
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
|
||||
):
|
||||
result = get_access_token()
|
||||
assert result is not None
|
||||
assert result.token == "valid-jwt"
|
||||
|
||||
async def test_token_without_expiry_always_valid(self):
|
||||
"""get_access_token() returns token when no expires_at is set."""
|
||||
from fastmcp.server.dependencies import _task_access_token
|
||||
|
||||
no_expiry = AccessToken(
|
||||
token="eternal-jwt",
|
||||
client_id="test-client",
|
||||
scopes=["read"],
|
||||
)
|
||||
_task_access_token.set(no_expiry)
|
||||
result = get_access_token()
|
||||
assert result is not None
|
||||
assert result.token == "eternal-jwt"
|
||||
_set_cached_snapshot(
|
||||
"test-task",
|
||||
TaskContextSnapshot(access_token_json=no_expiry.model_dump_json()),
|
||||
)
|
||||
fake_ctx = TaskContextInfo(task_id="test-task", session_id="s")
|
||||
with patch(
|
||||
"fastmcp.server.dependencies.get_task_context", return_value=fake_ctx
|
||||
):
|
||||
result = get_access_token()
|
||||
assert result is not None
|
||||
assert result.token == "eternal-jwt"
|
||||
|
||||
|
||||
class TestLifespanContextInBackgroundTasks:
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import pytest
|
|||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.dependencies import CurrentDocket, CurrentFastMCP, Depends
|
||||
from fastmcp.exceptions import ToolError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -257,8 +258,6 @@ async def test_dependency_errors_propagate_to_task_failure():
|
|||
) -> str:
|
||||
return f"Got: {dep}"
|
||||
|
||||
from fastmcp.exceptions import ToolError
|
||||
|
||||
async with Client(mcp) as client:
|
||||
task = await client.call_tool(
|
||||
"tool_with_failing_dep", {"value": "test"}, task=True
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ import asyncio
|
|||
import mcp.types as mt
|
||||
import pytest
|
||||
from docket import Docket
|
||||
from mcp.types import Tool as MCPTool
|
||||
from mcp.types import ToolExecution
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
|
|
@ -17,6 +19,7 @@ from fastmcp.prompts.base import PromptResult
|
|||
from fastmcp.resources.base import ResourceResult
|
||||
from fastmcp.server.dependencies import CurrentDocket, CurrentFastMCP
|
||||
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
||||
from fastmcp.server.providers.proxy import ProxyTool
|
||||
from fastmcp.server.tasks import TaskConfig
|
||||
from fastmcp.tools.base import ToolResult
|
||||
|
||||
|
|
@ -571,6 +574,21 @@ class TestMountedTaskMetadata:
|
|||
assert child_mcp_tool.execution.taskSupport == "optional"
|
||||
assert parent_mcp_tool.execution.taskSupport == "optional"
|
||||
|
||||
async def test_proxy_tool_preserves_execution_metadata(self):
|
||||
"""ProxyTool.from_mcp_tool should propagate execution.taskSupport (#3569)."""
|
||||
mcp_tool = MCPTool(
|
||||
name="remote_task_tool",
|
||||
description="A remote tool that supports tasks",
|
||||
inputSchema={"type": "object", "properties": {}},
|
||||
execution=ToolExecution(taskSupport="optional"),
|
||||
)
|
||||
|
||||
proxy = ProxyTool.from_mcp_tool(lambda: None, mcp_tool)
|
||||
result = proxy.to_mcp_tool(name=proxy.name)
|
||||
|
||||
assert result.execution is not None
|
||||
assert result.execution.taskSupport == "optional"
|
||||
|
||||
|
||||
class TestMountedTaskConfigModes:
|
||||
"""Test TaskConfig mode enforcement for mounted tools."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue