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

Restore proxy request context in all proxy forwarding handlers (fix v2 in-memory backend context leak)

Browse source
This commit is contained in:
Jeremiah Lowin 2026-07-05 23:29:59 -04:00
commit 7f3cee8243
No known key found for this signature in database
1 changed files with 82 additions and 83 deletions
Download patch file Download diff file Expand all files Collapse all files

165
fastmcp_slim/fastmcp/server/providers/proxy.py
View file

@ -85,7 +85,7 @@ class ProxyInitializeMiddleware(Middleware):
) -> mcp_types.InitializeResult | None:
client = await self.proxy._get_client()
try:
if isinstance(client, StatefulProxyClient):
if isinstance(client, ProxyClient):
ctx = context.fastmcp_context
if ctx is not None:
client._proxy_rc_ref[0] = (
@ -176,11 +176,12 @@ class ProxyTool(Tool):
client = await self._get_client()
async with client:
ctx = context or get_context()
# StatefulProxyClient reuses sessions across requests, so
# its receive-loop task has stale ContextVars from the first
# request. Stash the current RequestContext in the shared
# ref so handlers can restore it before forwarding.
if isinstance(client, StatefulProxyClient):
# Stash the proxy's RequestContext so the forwarding handlers can
# restore it before relaying a server-initiated request back to
# the proxy's client. Required for every proxy client: under SDK
# v2 an in-memory backend shares this event loop, so a handler's
# `get_context()` would otherwise resolve to the backend context.
if isinstance(client, ProxyClient):
client._proxy_rc_ref[0] = (
ctx.request_context,
ctx._fastmcp, # weakref to FastMCP, not the Context
@ -998,13 +999,17 @@ def _restore_request_context(
rc_ref: list[Any],
) -> None:
"""Set the ``request_ctx``, ``_current_context`` and ``_current_server``
ContextVars from stashed values.
ContextVars from stashed values so a proxy forwarding handler relays to the
proxy's own client rather than the upstream server.
Called at the start of proxy handler invocations in
``StatefulProxyClient`` to fix stale ContextVars in the receive-loop
task. Only overrides when the ContextVar is genuinely stale (same
session, different request_id) to avoid corrupting the concurrent
case where multiple sessions share the same ref via ``copy.copy``.
Called at the start of every proxy handler invocation. The stashed proxy
``RequestContext`` is the correct forwarding target, so we restore it unless
it is already active. This covers two cases:
- Stateful proxy: the reused receive-loop task carries a stale ContextVar
from an earlier request (same session, different request_id).
- In-memory backend (SDK v2): the backend runs in this event loop, so the
handler may inherit the *backend's* request_ctx (a different session).
We stash a ``(RequestContext, weakref[FastMCP])`` tuple — never a
``Context`` instance — because ``Context`` properties are themselves
@ -1031,19 +1036,14 @@ def _restore_request_context(
rc, fastmcp_ref = stashed
current_rc = fastmcp_request_ctx.get()
if current_rc is None:
fastmcp_request_ctx.set(rc)
fastmcp = fastmcp_ref()
if fastmcp is not None:
_current_context.set(Context(fastmcp))
_current_server.set(weakref.ref(fastmcp))
# Restore unless the stashed proxy context is already the active one.
if current_rc is rc:
return
if current_rc.session is rc.session and current_rc.request_id != rc.request_id:
fastmcp_request_ctx.set(rc)
fastmcp = fastmcp_ref()
if fastmcp is not None:
_current_context.set(Context(fastmcp))
_current_server.set(weakref.ref(fastmcp))
fastmcp_request_ctx.set(rc)
fastmcp = fastmcp_ref()
if fastmcp is not None:
_current_context.set(Context(fastmcp))
_current_server.set(weakref.ref(fastmcp))
def _make_restoring_handler(handler: Callable, rc_ref: list[Any]) -> Callable:
@ -1065,8 +1065,27 @@ class ProxyClient(Client[ClientTransportT]):
"""A proxy client that forwards advanced interactions between a remote MCP server and the proxy's connected clients.
Supports forwarding roots, sampling, elicitation, logging, and progress.
The default forwarding handlers must resolve the *proxy's* request context so
they relay server-initiated requests (roots/sampling/elicitation) back to the
proxy's own connected client, not to the upstream server they are talking to.
Under SDK v2 an in-memory backend runs in the same event loop as this client,
so a naive ``get_context()`` inside a handler can resolve to the backend's
context and forward the request straight back to the backend — an infinite
loop. To avoid that, ``ProxyTool.run`` (and the other proxy components) stash
the proxy-side ``RequestContext`` in ``_proxy_rc_ref`` before each backend
call, and the handlers are wrapped to restore it before forwarding.
"""
# Mutable list shared across copies (Client.new() uses copy.copy, which
# preserves references to mutable containers). Proxy components write [0]
# before each backend call; handlers read it to restore the proxy's
# request_ctx before forwarding. Stores a (RequestContext, weakref[FastMCP])
# tuple — never a Context instance — because Context properties are
# ContextVar-dependent and would resolve stale values in the receive loop.
_proxy_rc_ref: list[Any]
_proxy_restoring_handler_keys: set[str]
def __init__(
self,
transport: ClientTransportT
@ -1081,58 +1100,6 @@ class ProxyClient(Client[ClientTransportT]):
):
if "name" not in kwargs:
kwargs["name"] = self.generate_name()
if "roots" not in kwargs:
kwargs["roots"] = default_proxy_roots_handler
if "sampling_handler" not in kwargs:
kwargs["sampling_handler"] = default_proxy_sampling_handler
if "elicitation_handler" not in kwargs:
kwargs["elicitation_handler"] = default_proxy_elicitation_handler
if "log_handler" not in kwargs:
kwargs["log_handler"] = default_proxy_log_handler
if "progress_handler" not in kwargs:
kwargs["progress_handler"] = default_proxy_progress_handler
super().__init__(transport=transport, **kwargs)
# Enable forwarding of inbound HTTP headers (e.g. authorization) to
# the upstream server. This is only appropriate for proxy clients,
# where the caller's credentials should be propagated.
from fastmcp.client.transports.http import StreamableHttpTransport
from fastmcp.client.transports.sse import SSETransport
if isinstance(self.transport, StreamableHttpTransport | SSETransport):
self.transport.forward_incoming_headers = True
class StatefulProxyClient(ProxyClient[ClientTransportT]):
"""A proxy client that provides a stateful client factory for the proxy server.
The stateful proxy client bound its copy to the server session.
And it will be disconnected when the session is exited.
This is useful to proxy a stateful mcp server such as the Playwright MCP server.
Note that it is essential to ensure that the proxy server itself is also stateful.
Because session reuse means the receive-loop task inherits a stale
``request_ctx`` ContextVar snapshot, the default proxy handlers are
replaced with versions that restore the ContextVar before forwarding.
``ProxyTool.run`` stashes the current ``RequestContext`` in
``_proxy_rc_ref`` before each backend call, and the handlers consult
it to detect (and correct) staleness.
"""
# Mutable list shared across copies (Client.new() uses copy.copy,
# which preserves references to mutable containers). ProxyTool.run
# writes [0] before each backend call; handlers read it to detect
# stale ContextVars and restore the correct request_ctx.
#
# Stores a (RequestContext, weakref[FastMCP]) tuple — never a Context
# instance — because Context properties are ContextVar-dependent and
# would resolve stale values in the receive loop. The restore helper
# constructs a fresh Context from the weakref after setting request_ctx.
_proxy_rc_ref: list[Any]
_proxy_restoring_handler_keys: set[str]
def __init__(self, *args: Any, **kwargs: Any):
# Install context-restoring handler wrappers BEFORE super().__init__
# registers them with the Client's session kwargs.
self._proxy_rc_ref = [None]
@ -1147,12 +1114,16 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]):
if key not in kwargs:
kwargs[key] = _make_restoring_handler(default_fn, self._proxy_rc_ref)
self._proxy_restoring_handler_keys.add(key)
super().__init__(transport=transport, **kwargs)
super().__init__(*args, **kwargs)
# SDK v2 constructs a ServerSession per request, so per-session keying
# would build a fresh proxy client for every request. Key by the stable
# per-connection `Connection` instead, and tie cleanup to its exit stack.
self._caches: dict[Connection, Client[ClientTransportT]] = {}
# Enable forwarding of inbound HTTP headers (e.g. authorization) to
# the upstream server. This is only appropriate for proxy clients,
# where the caller's credentials should be propagated.
from fastmcp.client.transports.http import StreamableHttpTransport
from fastmcp.client.transports.sse import SSETransport
if isinstance(self.transport, StreamableHttpTransport | SSETransport):
self.transport.forward_incoming_headers = True
def _bind_restoring_handlers(self) -> None:
if "roots" in self._proxy_restoring_handler_keys:
@ -1180,12 +1151,40 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]):
default_proxy_progress_handler, self._proxy_rc_ref
)
def new(self) -> StatefulProxyClient[ClientTransportT]:
new_client = cast(StatefulProxyClient[ClientTransportT], super().new())
def new(self) -> ProxyClient[ClientTransportT]:
new_client = cast(ProxyClient[ClientTransportT], super().new())
new_client._proxy_rc_ref = [None]
new_client._proxy_restoring_handler_keys = set(
self._proxy_restoring_handler_keys
)
new_client._bind_restoring_handlers()
return new_client
class StatefulProxyClient(ProxyClient[ClientTransportT]):
"""A proxy client that provides a stateful client factory for the proxy server.
The stateful proxy client bound its copy to the server session.
And it will be disconnected when the session is exited.
This is useful to proxy a stateful mcp server such as the Playwright MCP server.
Note that it is essential to ensure that the proxy server itself is also stateful.
The base ``ProxyClient`` already installs the context-restoring handlers
(see its docstring); this subclass additionally caches one client per stable
``Connection`` and forces disconnect when the connection is torn down.
"""
def __init__(self, *args: Any, **kwargs: Any):
super().__init__(*args, **kwargs)
# SDK v2 constructs a ServerSession per request, so per-session keying
# would build a fresh proxy client for every request. Key by the stable
# per-connection `Connection` instead, and tie cleanup to its exit stack.
self._caches: dict[Connection, Client[ClientTransportT]] = {}
def new(self) -> StatefulProxyClient[ClientTransportT]:
return cast(StatefulProxyClient[ClientTransportT], super().new())
async def __aexit__(self, exc_type, exc_value, traceback) -> None: # type: ignore[override] # ty:ignore[invalid-method-override]
"""The stateful proxy client will be forced disconnected when the session is exited.