diff --git a/src/fastmcp/prompts/function_prompt.py b/src/fastmcp/prompts/function_prompt.py index b5ddd3425..a63167454 100644 --- a/src/fastmcp/prompts/function_prompt.py +++ b/src/fastmcp/prompts/function_prompt.py @@ -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]) diff --git a/src/fastmcp/resources/function_resource.py b/src/fastmcp/resources/function_resource.py index 771eeb7cb..b141a68bf 100644 --- a/src/fastmcp/resources/function_resource.py +++ b/src/fastmcp/resources/function_resource.py @@ -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]) diff --git a/src/fastmcp/resources/template.py b/src/fastmcp/resources/template.py index 3a3180c77..c9bc91f98 100644 --- a/src/fastmcp/resources/template.py +++ b/src/fastmcp/resources/template.py @@ -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]) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 11bbbf90e..7fe3bb62e 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -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 diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py index b024a1037..da8de3241 100644 --- a/src/fastmcp/server/dependencies.py +++ b/src/fastmcp/server/dependencies.py @@ -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: diff --git a/src/fastmcp/server/mixins/lifespan.py b/src/fastmcp/server/mixins/lifespan.py index a6c0cb9ed..508041009 100644 --- a/src/fastmcp/server/mixins/lifespan.py +++ b/src/fastmcp/server/mixins/lifespan.py @@ -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 diff --git a/src/fastmcp/server/providers/proxy.py b/src/fastmcp/server/providers/proxy.py index f2df251cf..ed8c61b94 100644 --- a/src/fastmcp/server/providers/proxy.py +++ b/src/fastmcp/server/providers/proxy.py @@ -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( diff --git a/src/fastmcp/server/tasks/handlers.py b/src/fastmcp/server/tasks/handlers.py index 051785ddd..2fe9eb83f 100644 --- a/src/fastmcp/server/tasks/handlers.py +++ b/src/fastmcp/server/tasks/handlers.py @@ -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 diff --git a/src/fastmcp/tools/function_tool.py b/src/fastmcp/tools/function_tool.py index 0f0dc8325..8dfc47f60 100644 --- a/src/fastmcp/tools/function_tool.py +++ b/src/fastmcp/tools/function_tool.py @@ -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 diff --git a/tests/server/tasks/test_concurrent_dependencies.py b/tests/server/tasks/test_concurrent_dependencies.py new file mode 100644 index 000000000..eb5af1bb1 --- /dev/null +++ b/tests/server/tasks/test_concurrent_dependencies.py @@ -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" diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py index 8b2a92889..ad472410f 100644 --- a/tests/server/tasks/test_context_background_task.py +++ b/tests/server/tasks/test_context_background_task.py @@ -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: diff --git a/tests/server/tasks/test_task_dependencies.py b/tests/server/tasks/test_task_dependencies.py index 7669f150c..129f0908d 100644 --- a/tests/server/tasks/test_task_dependencies.py +++ b/tests/server/tasks/test_task_dependencies.py @@ -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 diff --git a/tests/server/tasks/test_task_mount.py b/tests/server/tasks/test_task_mount.py index 1b1558579..03663b07f 100644 --- a/tests/server/tasks/test_task_mount.py +++ b/tests/server/tasks/test_task_mount.py @@ -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."""