From 44b2b17844baf49def5944b93094f7f8223bee51 Mon Sep 17 00:00:00 2001 From: William Easton Date: Wed, 13 Aug 2025 12:44:31 -0500 Subject: [PATCH] Fix mcp_config and simplify tests --- src/fastmcp/client/transports.py | 57 +++++++++----- src/fastmcp/mcp_config.py | 17 ++-- src/fastmcp/server/context.py | 11 ++- src/fastmcp/utilities/mcp_config.py | 73 ++++++++++++----- tests/client/test_stdio.py | 27 ++----- tests/test_mcp_config.py | 118 +++++++++++++++++----------- 6 files changed, 192 insertions(+), 111 deletions(-) diff --git a/src/fastmcp/client/transports.py b/src/fastmcp/client/transports.py index cf50848a3..bbf5ebfa0 100644 --- a/src/fastmcp/client/transports.py +++ b/src/fastmcp/client/transports.py @@ -298,11 +298,6 @@ class StreamableHttpTransport(ClientTransport): return f"" -class SessionHolder: - def __init__(self): - self.session: ClientSession | None = None - - class StdioTransport(ClientTransport): """ Base transport for connecting to an MCP server via subprocess with stdio. @@ -365,11 +360,11 @@ class StdioTransport(ClientTransport): if self._connect_task is not None: return - session_holder = SessionHolder() + session_future: asyncio.Future[ClientSession] = asyncio.Future() # start the connection task self._connect_task = asyncio.create_task( - _connect_task( + _stdio_transport_connect_task( command=self.command, args=self.args, env=self.env, @@ -377,7 +372,7 @@ class StdioTransport(ClientTransport): session_kwargs=session_kwargs, ready_event=self._ready_event, stop_event=self._stop_event, - session_holder=session_holder, + session_future=session_future, ) ) @@ -390,8 +385,8 @@ class StdioTransport(ClientTransport): if exception is not None: raise exception - self._session = session_holder.session - return session_holder.session + self._session = await session_future + return self._session async def disconnect(self): if self._connect_task is None: @@ -411,13 +406,18 @@ class StdioTransport(ClientTransport): async def close(self): await self.disconnect() + def __del__(self): + """Ensure that we send a disconnection signal to the transport task if we are being garbage collected.""" + if not self._stop_event.is_set(): + self._stop_event.set() + def __repr__(self) -> str: return ( f"<{self.__class__.__name__}(command='{self.command}', args={self.args})>" ) -async def _connect_task( +async def _stdio_transport_connect_task( command: str, args: list[str], env: dict[str, str] | None, @@ -425,8 +425,11 @@ async def _connect_task( session_kwargs: SessionKwargs, ready_event: anyio.Event, stop_event: anyio.Event, - session_holder: SessionHolder, + session_future: asyncio.Future[ClientSession], ): + """A standalone connection task for a stdio transport. It is not a part of the StdioTransport class + to ensure that the connection task does not hold a reference to the Transport object.""" + from mcp.client.stdio import stdio_client try: @@ -440,8 +443,10 @@ async def _connect_task( ) transport = await stack.enter_async_context(stdio_client(server_params)) read_stream, write_stream = transport - session_holder.session = await stack.enter_async_context( - ClientSession(read_stream, write_stream, **session_kwargs) + session_future.set_result( + await stack.enter_async_context( + ClientSession(read_stream, write_stream, **session_kwargs) + ) ) logger.debug("Stdio transport connected") @@ -451,7 +456,6 @@ async def _connect_task( await stop_event.wait() finally: # Clean up client on exit - session_holder.session = None logger.debug("Stdio transport disconnected") except Exception: # Ensure ready event is set even if connection fails @@ -860,12 +864,14 @@ class MCPConfigTransport(ClientTransport): """ def __init__(self, config: MCPConfig | dict, name_as_prefix: bool = True): - from fastmcp.utilities.mcp_config import composite_server_from_mcp_config + from fastmcp.utilities.mcp_config import mcp_config_to_servers_and_transports if isinstance(config, dict): config = MCPConfig.from_dict(config) self.config = config + self._underlying_transports: list[ClientTransport] = [] + # if there are no servers, raise an error if len(self.config.mcpServers) == 0: raise ValueError("No MCP servers defined in the config") @@ -873,14 +879,21 @@ class MCPConfigTransport(ClientTransport): # if there's exactly one server, create a client for that server elif len(self.config.mcpServers) == 1: self.transport = list(self.config.mcpServers.values())[0].to_transport() + self._underlying_transport = self.transport # otherwise create a composite client else: - self.transport = FastMCPTransport( - mcp=composite_server_from_mcp_config( - self.config, name_as_prefix=name_as_prefix + self._composite_server = FastMCP[Any]() + + for name, server, transport in mcp_config_to_servers_and_transports( + self.config + ): + self._underlying_transports.append(transport) + self._composite_server.mount( + server, prefix=name if name_as_prefix else None ) - ) + + self.transport = FastMCPTransport(mcp=self._composite_server) @contextlib.asynccontextmanager async def connect_session( @@ -889,6 +902,10 @@ class MCPConfigTransport(ClientTransport): async with self.transport.connect_session(**session_kwargs) as session: yield session + async def close(self): + for transport in self._underlying_transports: + await transport.close() + def __repr__(self) -> str: return f"" diff --git a/src/fastmcp/mcp_config.py b/src/fastmcp/mcp_config.py index a534fac38..5c2a54ded 100644 --- a/src/fastmcp/mcp_config.py +++ b/src/fastmcp/mcp_config.py @@ -47,11 +47,11 @@ from fastmcp.utilities.types import FastMCPBaseModel if TYPE_CHECKING: from fastmcp.client.transports import ( ClientTransport, - FastMCPTransport, SSETransport, StdioTransport, StreamableHttpTransport, ) + from fastmcp.server.server import FastMCP def infer_transport_type_from_url( @@ -90,10 +90,11 @@ class _TransformingMCPServerMixin(FastMCPBaseModel): description="The tags to exclude in the proxy.", ) - def to_transport(self) -> FastMCPTransport: - """Get the transport for the server.""" - from fastmcp.client.transports import FastMCPTransport - from fastmcp.server.server import FastMCP + def _to_server_and_transport( + self, + ) -> tuple[FastMCP[Any], ClientTransport]: + """Get the server and transport for the server.""" + from fastmcp import FastMCP transport: ClientTransport = super().to_transport() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue, reportUnknownVariableType] @@ -104,7 +105,11 @@ class _TransformingMCPServerMixin(FastMCPBaseModel): exclude_tags=self.exclude_tags, ) - return FastMCPTransport(wrapped_mcp_server) + return wrapped_mcp_server, transport + + def to_transport(self) -> ClientTransport: + """Get the transport for the server.""" + return self._to_server_and_transport()[1] class StdioMCPServer(BaseModel): diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 8a3901830..d5f238206 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import copy import warnings +import weakref from collections.abc import Generator, Mapping from contextlib import contextmanager from contextvars import ContextVar, Token @@ -115,11 +116,19 @@ class Context: """ def __init__(self, fastmcp: FastMCP): - self.fastmcp = fastmcp + self._fastmcp: weakref.ref[FastMCP] = weakref.ref(fastmcp) self._tokens: list[Token] = [] self._notification_queue: set[str] = set() # Dedupe notifications self._state: dict[str, Any] = {} + @property + def fastmcp(self) -> FastMCP: + """Get the FastMCP instance.""" + fastmcp = self._fastmcp() + if fastmcp is None: + raise RuntimeError("FastMCP instance is no longer available") + return fastmcp + async def __aenter__(self) -> Context: """Enter the context manager and set this context as the current context.""" parent_context = _current_context.get(None) diff --git a/src/fastmcp/utilities/mcp_config.py b/src/fastmcp/utilities/mcp_config.py index e648f4c1c..94467c5c4 100644 --- a/src/fastmcp/utilities/mcp_config.py +++ b/src/fastmcp/utilities/mcp_config.py @@ -1,28 +1,63 @@ from typing import Any -from fastmcp.mcp_config import MCPConfig +from fastmcp.client.transports import ClientTransport +from fastmcp.mcp_config import ( + MCPConfig, + MCPServerTypes, +) from fastmcp.server.server import FastMCP +# def composite_server_from_mcp_config( +# config: MCPConfig, name_as_prefix: bool = True +# ) -> tuple[FastMCP[None], list[ClientTransport]]: +# """A utility function to create a composite server from an MCPConfig, returns the underlying +# transports for each server. +# """ +# composite_server = FastMCP[None]() -def composite_server_from_mcp_config( - config: MCPConfig, name_as_prefix: bool = True -) -> FastMCP[None]: - """A utility function to create a composite server from an MCPConfig.""" - composite_server = FastMCP[None]() +# transports = mount_mcp_config_into_server(config, composite_server, name_as_prefix) - mount_mcp_config_into_server(config, composite_server, name_as_prefix) - - return composite_server +# return composite_server, transports -def mount_mcp_config_into_server( +# def mount_mcp_config_into_server( +# config: MCPConfig, +# server: FastMCP[Any], +# name_as_prefix: bool = True, +# ) -> None: +# """A utility function to mount the servers from an MCPConfig into a FastMCP server, returns the underlying +# transports for each server. +# """ +# for name, server_to_mount, transport in mcp_config_to_servers_and_transports(config): +# server.mount(server=server_to_mount, prefix=name if name_as_prefix else None) + +def mcp_config_to_servers_and_transports( config: MCPConfig, - server: FastMCP[Any], - name_as_prefix: bool = True, -) -> None: - """A utility function to mount the servers from an MCPConfig into a FastMCP server.""" - for name, mcp_server in config.mcpServers.items(): - server.mount( - prefix=name if name_as_prefix else None, - server=FastMCP.as_proxy(backend=mcp_server.to_transport()), - ) +) -> list[tuple[str, FastMCP[Any], ClientTransport]]: + """A utility function to convert each entry of an MCP Config into a transport and server.""" + return [ + mcp_server_type_to_servers_and_transports(name, mcp_server) + for name, mcp_server in config.mcpServers.items() + ] + +def mcp_server_type_to_servers_and_transports( + name: str, + mcp_server: MCPServerTypes, +) -> tuple[str, FastMCP[Any], ClientTransport]: + """A utility function to convert each entry of an MCP Config into a transport and server.""" + from fastmcp.mcp_config import ( + TransformingRemoteMCPServer, + TransformingStdioMCPServer, +) + server: FastMCP[Any] + transport: ClientTransport + + if isinstance( + mcp_server, TransformingRemoteMCPServer | TransformingStdioMCPServer + ): + server, transport = mcp_server._to_server_and_transport() + else: + transport = mcp_server.to_transport() + server = FastMCP.as_proxy(backend=transport) + + return name, server, transport diff --git a/tests/client/test_stdio.py b/tests/client/test_stdio.py index 05fb62b83..cbac8740e 100644 --- a/tests/client/test_stdio.py +++ b/tests/client/test_stdio.py @@ -83,6 +83,8 @@ class TestKeepAlive: gc_collect_harder() + # When debugging, the debugger holds extra references so this test + # will pass when running and fail under the debugger. assert transport_weak_ref transport = transport_weak_ref() assert transport is None @@ -90,17 +92,10 @@ class TestKeepAlive: async def test_keep_alive_true_exit_scope_kills_client(self, stdio_script): pid: int | None = None - transport_weak_ref: weakref.ref[PythonStdioTransport] | None = None - client_weak_ref: weakref.ref[Client] | None = None - async def test_server(): transport = PythonStdioTransport(script_path=stdio_script, keep_alive=True) client = Client(transport=transport) - nonlocal client_weak_ref - client_weak_ref = weakref.ref(client) - nonlocal transport_weak_ref - transport_weak_ref = weakref.ref(transport) assert client.transport.keep_alive is True async with client: @@ -112,18 +107,10 @@ class TestKeepAlive: gc_collect_harder() - await asyncio.sleep(1) - - assert client_weak_ref - client = client_weak_ref() - assert client is None - - assert transport_weak_ref - transport = transport_weak_ref() - assert transport is None - with pytest.raises(psutil.NoSuchProcess): - psutil.Process(pid) + while True: + psutil.Process(pid) + await asyncio.sleep(0.1) async def test_keep_alive_false_exit_scope_kills_server(self, stdio_script): pid: int | None = None @@ -142,7 +129,9 @@ class TestKeepAlive: await test_server() with pytest.raises(psutil.NoSuchProcess): - psutil.Process(pid) + while True: + psutil.Process(pid) + await asyncio.sleep(0.1) async def test_keep_alive_false_starts_new_session_across_multiple_calls( self, stdio_script diff --git a/tests/test_mcp_config.py b/tests/test_mcp_config.py index 9b77f3fd9..961140ee5 100644 --- a/tests/test_mcp_config.py +++ b/tests/test_mcp_config.py @@ -3,7 +3,6 @@ import gc import inspect import logging import tempfile -import weakref from collections.abc import AsyncGenerator from pathlib import Path from typing import Any @@ -16,8 +15,6 @@ from fastmcp.client.auth.oauth import OAuthClientProvider from fastmcp.client.client import Client from fastmcp.client.logging import LogMessage from fastmcp.client.transports import ( - ClientTransport, - FastMCPTransport, MCPConfigTransport, SSETransport, StdioTransport, @@ -32,7 +29,6 @@ from fastmcp.mcp_config import ( StdioMCPServer, TransformingStdioMCPServer, ) -from fastmcp.server.server import FastMCP from fastmcp.tools.tool import Tool as FastMCPTool @@ -242,6 +238,64 @@ async def test_multi_client(tmp_path: Path): async def test_multi_client_lifespan(tmp_path: Path): + pid_1: int | None = None + pid_2: int | None = None + + async def test_server(): + server_script = inspect.cleandoc(""" + from fastmcp import FastMCP + import os + + mcp = FastMCP() + + @mcp.tool + def pid() -> int: + return os.getpid() + + if __name__ == '__main__': + mcp.run() + """) + + script_path = tmp_path / "test.py" + script_path.write_text(server_script) + + config = { + "mcpServers": { + "test_1": { + "command": "python", + "args": [str(script_path)], + }, + "test_2": { + "command": "python", + "args": [str(script_path)], + }, + } + } + transport = MCPConfigTransport(config) + client = Client(transport) + + async with client: + nonlocal pid_1 + pid_1 = (await client.call_tool("test_1_pid")).data + + nonlocal pid_2 + pid_2 = (await client.call_tool("test_2_pid")).data + + await test_server() + + gc_collect_harder() + + with pytest.raises(psutil.NoSuchProcess): + while True: + psutil.Process(pid_1) + await asyncio.sleep(0.1) + + with pytest.raises(psutil.NoSuchProcess): + while True: + psutil.Process(pid_2) + await asyncio.sleep(0.1) + +async def test_multi_client_force_close(tmp_path: Path): server_script = inspect.cleandoc(""" from fastmcp import FastMCP import os @@ -271,54 +325,26 @@ async def test_multi_client_lifespan(tmp_path: Path): }, } } + transport = MCPConfigTransport(config) + client = Client(transport) - pid: int | None = None + async with client: + pid_1 = (await client.call_tool("test_1_pid")).data + pid_2 = (await client.call_tool("test_2_pid")).data - transport_weak_ref: weakref.ref[ClientTransport] | None = None - nested_transport_weak_ref: weakref.ref[ClientTransport] | None = None - client_weak_ref: weakref.ref[Client] | None = None - server_weak_ref: weakref.ref[FastMCP] | None = None - - async def test_server(): - transport = MCPConfigTransport(config) - client = Client(transport) - nonlocal client_weak_ref - client_weak_ref = weakref.ref(client) - nonlocal transport_weak_ref - transport_weak_ref = weakref.ref(transport) - - nonlocal nested_transport_weak_ref - nested_transport_weak_ref = weakref.ref(transport.transport) - assert isinstance(transport.transport, FastMCPTransport) - nonlocal server_weak_ref - server_weak_ref = weakref.ref(transport.transport.server) - - async with client: - nonlocal pid - pid = (await client.call_tool("test_1_pid")).data - - await test_server() + await client.close() gc_collect_harder() - await asyncio.sleep(1) - - gc_collect_harder() - - await asyncio.sleep(1) - - assert client_weak_ref is not None - assert transport_weak_ref is not None - assert nested_transport_weak_ref is not None - assert server_weak_ref is not None - - assert not client_weak_ref() - assert not transport_weak_ref() - assert not nested_transport_weak_ref() - assert not server_weak_ref() - with pytest.raises(psutil.NoSuchProcess): - psutil.Process(pid) + process = psutil.Process(pid_1) + + assert not process + + with pytest.raises(psutil.NoSuchProcess): + process = psutil.Process(pid_2) + + assert not process async def test_remote_config_default_no_auth():