diff --git a/src/fastmcp/client/transports/config.py b/src/fastmcp/client/transports/config.py index 5f303567d..c8d5a8a40 100644 --- a/src/fastmcp/client/transports/config.py +++ b/src/fastmcp/client/transports/config.py @@ -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: diff --git a/src/fastmcp/server/providers/proxy.py b/src/fastmcp/server/providers/proxy.py index be844bb4a..5d2ea562f 100644 --- a/src/fastmcp/server/providers/proxy.py +++ b/src/fastmcp/server/providers/proxy.py @@ -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): diff --git a/tests/test_mcp_config.py b/tests/test_mcp_config.py index 51eeae5e1..786d4e4e2 100644 --- a/tests/test_mcp_config.py +++ b/tests/test_mcp_config.py @@ -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():