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

Fix mcp_config and simplify tests

Browse source
This commit is contained in:
William Easton 2025-08-13 12:44:31 -05:00
commit 44b2b17844
No known key found for this signature in database
6 changed files with 189 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

57
src/fastmcp/client/transports.py
View file

@ -298,11 +298,6 @@ class StreamableHttpTransport(ClientTransport):
return f"<StreamableHttpTransport(url='{self.url}')>"
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"<MCPConfigTransport(config='{self.config}')>"

17
src/fastmcp/mcp_config.py
View file

@ -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):

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

@ -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)

73
src/fastmcp/utilities/mcp_config.py
View file

@ -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

27
tests/client/test_stdio.py
View file

@ -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

118
tests/test_mcp_config.py
View file

@ -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():