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

Fix session persistence across tool calls in multi-server MCPConfigTransport (#3330)

Browse source 
* connect ProxyClient via AsyncExitStack

* test session persistence for multiple tool calls

* switch to StatefulProxyClient

* changed to StatefulProxyClient

* stash weakref instead of context

* clear transport list before building new session
This commit is contained in:
Jeremy Simon 2026-03-02 23:19:43 +01:00 • committed by GitHub
commit f2dd2e5124
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 157 additions and 58 deletions
Download patch file Download diff file Expand all files Collapse all files

79
src/fastmcp/client/transports/config.py
View file

@ -91,40 +91,57 @@ class MCPConfigTransport(ClientTransport):
yield session
return
# Multiple servers - create composite with mounted proxies
# Close any previous transports from prior connections to avoid leaking
for t in self._transports:
await t.close()
self._transports = []
# Multiple servers - create composite with mounted proxies, connecting
# each ProxyClient so its underlying transport session stays alive for
# the duration of this context (fixes session persistence for
# streamable-http backends — see #2790).
timeout = session_kwargs.get("read_timeout_seconds")
composite = FastMCP[Any](name="MCPRouter")
try:
for name, server_config in self.config.mcpServers.items():
transport, proxy = self._create_proxy(name, server_config, timeout)
self._transports.append(transport)
composite.mount(proxy, namespace=name if self.name_as_prefix else None)
except Exception:
# Clean up any transports created before the failure
async with contextlib.AsyncExitStack() as stack:
# Close any previous transports from prior connections to avoid leaking
for t in self._transports:
await t.close()
self._transports = []
raise
async with FastMCPTransport(mcp=composite).connect_session(
**session_kwargs
) as session:
yield session
try:
for name, server_config in self.config.mcpServers.items():
transport, _client, proxy = await self._create_proxy(
name, server_config, timeout, stack
)
self._transports.append(transport)
composite.mount(
proxy, namespace=name if self.name_as_prefix else None
)
except Exception:
# Clean up any transports created before the failure
for t in self._transports:
await t.close()
self._transports = []
raise
def _create_proxy(
async with FastMCPTransport(mcp=composite).connect_session(
**session_kwargs
) as session:
yield session
async def _create_proxy(
self,
name: str,
config: MCPServerTypes,
timeout: datetime.timedelta | None,
) -> tuple[ClientTransport, FastMCP[Any]]:
"""Create underlying transport and proxy server for a single backend."""
stack: contextlib.AsyncExitStack,
) -> tuple[ClientTransport, Any, FastMCP[Any]]:
"""Create underlying transport, proxy client, and proxy server for a single backend.
The ProxyClient is connected via the AsyncExitStack *before* being
passed to create_proxy so the factory sees it as connected and reuses
the same session for all tool calls (instead of creating fresh copies).
Returns a tuple of (transport, proxy_client, proxy_server).
"""
# Import here to avoid circular dependency
from fastmcp.server.providers.proxy import ProxyClient
from fastmcp.server.providers.proxy import StatefulProxyClient
tool_transforms = None
include_tags = None
@ -144,7 +161,23 @@ class MCPConfigTransport(ClientTransport):
else:
transport = config.to_transport()
client = ProxyClient(transport=transport, timeout=timeout)
client = StatefulProxyClient(transport=transport, timeout=timeout)
# Connect the client *before* create_proxy so _create_client_factory
# detects it as connected and reuses it for all tool calls, preserving
# the session ID across requests. StatefulProxyClient is used instead
# of ProxyClient because its context-restoring handler wrappers prevent
# stale ContextVars in the reused session's receive loop.
#
# StatefulProxyClient.__aexit__ is a no-op (by design, for the
# new_stateful() use case), so we cannot rely on enter_async_context
# alone to clean up. Instead we connect manually and push an
# explicit force-disconnect callback so the subprocess is terminated
# when the AsyncExitStack unwinds.
await client.__aenter__()
# Callbacks run LIFO: transport.close() must run *after*
# client._disconnect so push it first.
stack.push_async_callback(transport.close)
stack.push_async_callback(client._disconnect, force=True)
# Create proxy without include_tags/exclude_tags - we'll add them after tool transforms
proxy = create_proxy(
client,
@ -160,7 +193,7 @@ class MCPConfigTransport(ClientTransport):
proxy.enable(tags=set(include_tags), only=True)
if exclude_tags:
proxy.disable(tags=set(exclude_tags))
return transport, proxy
return transport, client, proxy
async def close(self):
for transport in self._transports:

34
src/fastmcp/server/providers/proxy.py
View file

@ -127,7 +127,10 @@ class ProxyTool(Tool):
# request. Stash the current RequestContext in the shared
# ref so handlers can restore it before forwarding.
if isinstance(client, StatefulProxyClient):
cast(list[Any], client._proxy_rc_ref)[0] = ctx.request_context
cast(list[Any], client._proxy_rc_ref)[0] = (
ctx.request_context,
ctx._fastmcp, # weakref to FastMCP, not the Context
)
# Build meta dict from request context
meta: dict[str, Any] | None = None
if hasattr(ctx, "request_context"):
@ -791,24 +794,40 @@ async def default_proxy_progress_handler(
def _restore_request_context(
rc_ref: list[Any],
) -> None:
"""Set the ``request_ctx`` ContextVar from a stashed RequestContext.
"""Set the ``request_ctx`` and ``_current_context`` ContextVars from stashed values.
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``.
We stash a ``(RequestContext, weakref[FastMCP])`` tuple — never a
``Context`` instance — because ``Context`` properties are themselves
ContextVar-dependent and would resolve stale values in the receive
loop. Instead we construct a fresh ``Context`` here after restoring
``request_ctx``, so its property accesses read the correct values.
"""
rc = rc_ref[0]
if rc is None:
from fastmcp.server.context import Context, _current_context
stashed = rc_ref[0]
if stashed is None:
return
rc, fastmcp_ref = stashed
try:
current_rc = request_ctx.get()
except LookupError:
request_ctx.set(rc)
fastmcp = fastmcp_ref()
if fastmcp is not None:
_current_context.set(Context(fastmcp))
return
if current_rc.session is rc.session and current_rc.request_id != rc.request_id:
request_ctx.set(rc)
fastmcp = fastmcp_ref()
if fastmcp is not None:
_current_context.set(Context(fastmcp))
def _make_restoring_handler(handler: Callable, rc_ref: list[Any]) -> Callable:
@ -881,9 +900,10 @@ class StatefulProxyClient(ProxyClient[ClientTransportT]):
# writes [0] before each backend call; handlers read it to detect
# stale ContextVars and restore the correct request_ctx.
#
# We store the concrete RequestContext (not fastmcp's Context) because
# Context properties are themselves ContextVar-dependent and resolve
# in the caller's async context — which is stale in the receive loop.
# 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]
def __init__(self, *args: Any, **kwargs: Any):

100
tests/test_mcp_config.py
View file

@ -902,37 +902,83 @@ async def test_multi_server_timeout_propagation():
transport = MCPConfigTransport(config)
timeout = timedelta(seconds=42)
# Patch _create_proxy to verify timeout is passed correctly
# Mock _create_proxy to avoid real stdio connections and verify timeout
mock_create_proxy = AsyncMock(
return_value=(AsyncMock(), AsyncMock(), FastMCP(name="MockProxy"))
)
with (
patch("fastmcp.client.transports.FastMCP.as_proxy") as mock_as_proxy,
patch.object(
transport, "_create_proxy", wraps=transport._create_proxy
) as mock_create_proxy,
):
# Make as_proxy return a mock FastMCP
mock_proxy = FastMCP(name="MockProxy")
mock_as_proxy.return_value = mock_proxy
# Mock connect_session on FastMCPTransport to avoid actual connection
with patch(
patch.object(transport, "_create_proxy", mock_create_proxy),
patch(
"fastmcp.client.transports.FastMCPTransport.connect_session"
) as mock_connect:
mock_session = AsyncMock()
mock_connect.return_value.__aenter__ = AsyncMock(return_value=mock_session)
mock_connect.return_value.__aexit__ = AsyncMock(return_value=None)
) as mock_connect,
):
mock_session = AsyncMock()
mock_connect.return_value.__aenter__ = AsyncMock(return_value=mock_session)
mock_connect.return_value.__aexit__ = AsyncMock(return_value=None)
async with transport.connect_session(read_timeout_seconds=timeout):
pass
async with transport.connect_session(read_timeout_seconds=timeout):
pass
# Verify _create_proxy was called with the timeout for each server
assert mock_create_proxy.call_count == 2
for call in mock_create_proxy.call_args_list:
_, kwargs = call.args, call.kwargs if call.kwargs else {}
# Third positional arg is timeout
call_timeout = call[0][2] if len(call[0]) > 2 else kwargs.get("timeout")
assert call_timeout == timeout, (
f"Expected timeout {timeout}, got {call_timeout}"
)
# Verify _create_proxy was called with the timeout for each server
assert mock_create_proxy.call_count == 2
for call in mock_create_proxy.call_args_list:
# Third positional arg is timeout
call_timeout = call[0][2] if len(call[0]) > 2 else call.kwargs.get("timeout")
assert call_timeout == timeout, (
f"Expected timeout {timeout}, got {call_timeout}"
)
async def test_multi_server_session_persistence(tmp_path: Path):
"""Test that session IDs persist across tool calls in multi-server mode.
Regression test for https://github.com/PrefectHQ/fastmcp/issues/2790 —
MCPConfigTransport was not connecting ProxyClients before mounting, so
each tool call opened a new session with the backend server.
"""
server_script = inspect.cleandoc("""
from fastmcp import FastMCP, Context
mcp = FastMCP()
@mcp.tool
def get_session(ctx: Context) -> str:
return ctx.session_id
if __name__ == '__main__':
mcp.run()
""")
script_path = tmp_path / "session_server.py"
script_path.write_text(server_script)
config = {
"mcpServers": {
"server1": {
"command": "python",
"args": [str(script_path)],
},
"server2": {
"command": "python",
"args": [str(script_path)],
},
}
}
client = Client(config)
async with client:
result1 = await client.call_tool("server1_get_session", {})
assert isinstance(result1.content[0], TextContent)
session_id_1 = result1.content[0].text
result2 = await client.call_tool("server1_get_session", {})
assert isinstance(result2.content[0], TextContent)
session_id_2 = result2.content[0].text
assert session_id_1 == session_id_2, (
f"Session ID changed between calls: {session_id_1} != {session_id_2}"
)
async def test_single_server_config_transport():