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

Merge pull request #2844 from jlowin/fix-asgi-contextvar-propagation

Browse source 
Fix ContextVar propagation for ASGI-mounted servers with tasks
This commit is contained in:
Chris Guidry 2026-01-12 10:38:31 -05:00 • committed by GitHub
commit 4ae0cc205c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 33 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

32
src/fastmcp/server/context.py
View file

@ -185,10 +185,24 @@ class Context:
self._tokens.append(token)
# Set current server for dependency injection (use weakref to avoid reference cycles)
from fastmcp.server.dependencies import _current_server
from fastmcp.server.dependencies import (
_current_docket,
_current_server,
_current_worker,
)
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
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)
# Start background notification flusher
self._exit_stack = AsyncExitStack()
await self._exit_stack.__aenter__()
@ -207,10 +221,20 @@ class Context:
if self._exit_stack is not None:
await self._exit_stack.aclose()
# Reset server token
if hasattr(self, "_server_token"):
from fastmcp.server.dependencies import _current_server
# Reset server/docket/worker tokens
from fastmcp.server.dependencies import (
_current_docket,
_current_server,
_current_worker,
)
if hasattr(self, "_worker_token"):
_current_worker.reset(self._worker_token)
delattr(self, "_worker_token")
if hasattr(self, "_docket_token"):
_current_docket.reset(self._docket_token)
delattr(self, "_docket_token")
if hasattr(self, "_server_token"):
_current_server.reset(self._server_token)
delattr(self, "_server_token")

6
src/fastmcp/server/server.py
View file

@ -245,8 +245,9 @@ class FastMCP(Generic[LifespanResultT]):
# Resolve server default for background task support
self._support_tasks_by_default: bool = tasks if tasks is not None else False
# Docket instance (set during lifespan for cross-task access)
# Docket and Worker instances (set during lifespan for cross-task access)
self._docket = None
self._worker = None
self._additional_http_routes: list[BaseRoute] = []
@ -538,6 +539,8 @@ class FastMCP(Generic[LifespanResultT]):
# Create and start Worker
async with Worker(docket, **worker_kwargs) as worker: # type: ignore[arg-type]
# 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:
@ -550,6 +553,7 @@ class FastMCP(Generic[LifespanResultT]):
await worker_task
finally:
_current_worker.reset(worker_token)
self._worker = None
finally:
# Reset ContextVar
_current_docket.reset(docket_token)