Add keep_alive param to reuse subprocess

This commit is contained in:
Jeremiah Lowin 2025-05-28 21:46:40 -04:00
commit 9bbd417129
5 changed files with 366 additions and 28 deletions

View file

@ -18,18 +18,6 @@ The FastMCP Client architecture separates the protocol logic (`Client`) from the
- **`Client`**: Handles sending MCP requests (like `tools/call`, `resources/read`), receiving responses, and managing callbacks.
- **`Transport`**: Responsible for establishing and maintaining the connection to the server (e.g., via WebSockets, SSE, Stdio, or in-memory).
```python
from fastmcp import Client, FastMCP
from fastmcp.client import (
RootsHandler,
RootsList,
LogHandler,
MessageHandler,
SamplingHandler,
ProgressHandler # For handling progress notifications
)
```
### Transports
Clients must be initialized with a `transport`. You can either provide an already instantiated transport object, or provide a transport source and let FastMCP attempt to infer the correct transport to use.
@ -282,7 +270,7 @@ These methods are especially useful for debugging or when you need to access met
### Additional Features
#### Pinging the server
#### Pinging the Server
The client can be used to ping the server to verify connectivity.
@ -292,6 +280,62 @@ async with client:
print("Server is reachable")
```
#### Session Management
When using stdio transports, clients support a `keep_alive` feature (enabled by default) that maintains subprocess sessions between connection contexts. You can manually control this behavior using the client's `close()` method.
When `keep_alive=False`, the client will automatically close the session when the context manager exits.
<CodeGroup>
```python keep_alive=True
from fastmcp import Client
# Client with keep_alive=True (default)
client = Client("my_mcp_server.py")
async def example():
# First session
async with client:
await client.ping()
# Second session - uses the same subprocess
async with client:
await client.ping()
# Manually close the session
await client.close()
# Third session - will start a new subprocess
async with client:
await client.ping()
asyncio.run(example())
```
```python keep_alive=False
from fastmcp import Client
# Client with keep_alive=False
client = Client("my_mcp_server.py", keep_alive=False)
async def example():
# First session
async with client:
await client.ping()
# Second session - will start a new subprocess
async with client:
await client.ping()
# Third session - will start a new subprocess
async with client:
await client.ping()
asyncio.run(example())
```
</CodeGroup>
#### Timeouts
<VersionBadge version="2.3.4" />

View file

@ -160,6 +160,63 @@ client = Client(transport)
These transports manage an MCP server running as a subprocess, communicating with it via standard input (stdin) and standard output (stdout). This is the standard mechanism used by clients like Claude Desktop.
### Session Management
All stdio transports support a `keep_alive` parameter (default: `True`) that controls session persistence across multiple client context managers:
- **`keep_alive=True` (default)**: The subprocess and session are maintained between client context exits and re-entries. This improves performance when making multiple separate connections to the same server.
- **`keep_alive=False`**: A new subprocess is started for each client context, ensuring complete isolation between sessions.
When `keep_alive=True`, you can manually close the session using `await client.close()` if needed. This will terminate the subprocess and require a new one to be started on the next connection.
<CodeGroup>
```python keep_alive=True
from fastmcp import Client
# Client with keep_alive=True (default)
client = Client("my_mcp_server.py")
async def example():
# First session
async with client:
await client.ping()
# Second session - uses the same subprocess
async with client:
await client.ping()
# Manually close the session
await client.close()
# Third session - will start a new subprocess
async with client:
await client.ping()
asyncio.run(example())
```
```python keep_alive=False
from fastmcp import Client
# Client with keep_alive=False
client = Client("my_mcp_server.py", keep_alive=False)
async def example():
# First session
async with client:
await client.ping()
# Second session - will start a new subprocess
async with client:
await client.ping()
# Third session - will start a new subprocess
async with client:
await client.ping()
asyncio.run(example())
```
</CodeGroup>
### Python Stdio
- **Class:** `fastmcp.client.transports.PythonStdioTransport`
@ -218,7 +275,7 @@ client = Client(node_server_script)
# Option 2: Explicit transport
transport = NodeStdioTransport(
script_path=node_server_script,
node_cmd="node" # Optional: specify path to Node executable
node_cmd="node", # Optional: specify path to Node executable
)
client = Client(transport)

View file

@ -194,6 +194,7 @@ class Client(Generic[ClientTransportT]):
raise RuntimeError(
"Client is not connected. Use the 'async with client:' context manager first."
)
return self._session
@property
@ -231,6 +232,8 @@ class Client(Generic[ClientTransportT]):
with anyio.fail_after(self._init_timeout):
self._initialize_result = await self._session.initialize()
yield
except anyio.ClosedResourceError:
raise RuntimeError("Server session was closed unexpectedly")
except TimeoutError:
raise RuntimeError("Failed to initialize server session")
finally:
@ -263,6 +266,9 @@ class Client(Generic[ClientTransportT]):
finally:
self._exit_stack = None
async def close(self):
await self.transport.close()
# --- MCP Client Methods ---
async def ping(self) -> bool:

View file

@ -1,9 +1,11 @@
import abc
import asyncio
import contextlib
import datetime
import os
import shutil
import sys
import warnings
from collections.abc import AsyncIterator
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypedDict, TypeVar, cast, overload
@ -86,11 +88,20 @@ class ClientTransport(abc.ABC):
# Basic representation for subclasses
return f"<{self.__class__.__name__}>"
async def close(self):
"""Close the transport."""
pass
class WSTransport(ClientTransport):
"""Transport implementation that connects to an MCP server via WebSockets."""
def __init__(self, url: str | AnyUrl):
# we never really used this transport, so it can be removed at any time
warnings.warn(
"WSTransport is a deprecated MCP transport and will be removed in a future version. Use StreamableHttpTransport instead.",
DeprecationWarning,
)
if isinstance(url, AnyUrl):
url = str(url)
if not isinstance(url, str) or not url.startswith("ws"):
@ -227,6 +238,7 @@ class StdioTransport(ClientTransport):
args: list[str],
env: dict[str, str] | None = None,
cwd: str | None = None,
keep_alive: bool | None = None,
):
"""
Initialize a Stdio transport.
@ -241,20 +253,81 @@ class StdioTransport(ClientTransport):
self.args = args
self.env = env
self.cwd = cwd
if keep_alive is None:
keep_alive = True
self.keep_alive = keep_alive
self._session: ClientSession | None = None
self._connect_task: asyncio.Task | None = None
self._ready_event = asyncio.Event()
self._stop_event = asyncio.Event()
@contextlib.asynccontextmanager
async def connect_session(
self, **session_kwargs: Unpack[SessionKwargs]
) -> AsyncIterator[ClientSession]:
server_params = StdioServerParameters(
command=self.command, args=self.args, env=self.env, cwd=self.cwd
)
async with stdio_client(server_params) as transport:
read_stream, write_stream = transport
async with ClientSession(
read_stream, write_stream, **session_kwargs
) as session:
yield session
try:
await self.connect(**session_kwargs)
assert self._session is not None
yield self._session
finally:
if not self.keep_alive:
await self.disconnect()
else:
logger.debug("Stdio transport has keep_alive=True, not disconnecting")
async def connect(
self, **session_kwargs: Unpack[SessionKwargs]
) -> ClientSession | None:
if self._connect_task is not None:
return
async def _connect_task():
async with contextlib.AsyncExitStack() as stack:
try:
server_params = StdioServerParameters(
command=self.command, args=self.args, env=self.env, cwd=self.cwd
)
transport = await stack.enter_async_context(
stdio_client(server_params)
)
read_stream, write_stream = transport
self._session = await stack.enter_async_context(
ClientSession(read_stream, write_stream, **session_kwargs)
)
logger.debug("Stdio transport connected")
self._ready_event.set()
# Wait until disconnect is requested (stop_event is set)
await self._stop_event.wait()
finally:
# Clean up client on exit
self._session = None
logger.debug("Stdio transport disconnected")
# start the connection task
self._connect_task = asyncio.create_task(_connect_task())
# wait for the client to be ready before returning
await self._ready_event.wait()
async def disconnect(self):
if self._connect_task is None:
return
# signal the connection task to stop
self._stop_event.set()
# wait for the connection task to finish cleanly
await self._connect_task
# reset variables and events for potential future reconnects
self._connect_task = None
self._stop_event = asyncio.Event()
self._ready_event = asyncio.Event()
async def close(self):
await self.disconnect()
def __repr__(self) -> str:
return (
@ -272,6 +345,7 @@ class PythonStdioTransport(StdioTransport):
env: dict[str, str] | None = None,
cwd: str | None = None,
python_cmd: str = sys.executable,
keep_alive: bool | None = None,
):
"""
Initialize a Python transport.
@ -293,7 +367,13 @@ class PythonStdioTransport(StdioTransport):
if args:
full_args.extend(args)
super().__init__(command=python_cmd, args=full_args, env=env, cwd=cwd)
super().__init__(
command=python_cmd,
args=full_args,
env=env,
cwd=cwd,
keep_alive=keep_alive,
)
self.script_path = script_path
@ -306,6 +386,7 @@ class FastMCPStdioTransport(StdioTransport):
args: list[str] | None = None,
env: dict[str, str] | None = None,
cwd: str | None = None,
keep_alive: bool | None = None,
):
script_path = Path(script_path).resolve()
if not script_path.is_file():
@ -314,7 +395,11 @@ class FastMCPStdioTransport(StdioTransport):
raise ValueError(f"Not a Python script: {script_path}")
super().__init__(
command="fastmcp", args=["run", str(script_path)], env=env, cwd=cwd
command="fastmcp",
args=["run", str(script_path)],
env=env,
cwd=cwd,
keep_alive=keep_alive,
)
self.script_path = script_path
@ -329,6 +414,7 @@ class NodeStdioTransport(StdioTransport):
env: dict[str, str] | None = None,
cwd: str | None = None,
node_cmd: str = "node",
keep_alive: bool | None = None,
):
"""
Initialize a Node transport.
@ -350,7 +436,9 @@ class NodeStdioTransport(StdioTransport):
if args:
full_args.extend(args)
super().__init__(command=node_cmd, args=full_args, env=env, cwd=cwd)
super().__init__(
command=node_cmd, args=full_args, env=env, cwd=cwd, keep_alive=keep_alive
)
self.script_path = script_path
@ -366,6 +454,7 @@ class UvxStdioTransport(StdioTransport):
with_packages: list[str] | None = None,
from_package: str | None = None,
env_vars: dict[str, str] | None = None,
keep_alive: bool | None = None,
):
"""
Initialize a Uvx transport.
@ -405,7 +494,13 @@ class UvxStdioTransport(StdioTransport):
env = os.environ.copy()
env.update(env_vars)
super().__init__(command="uvx", args=uvx_args, env=env, cwd=project_directory)
super().__init__(
command="uvx",
args=uvx_args,
env=env,
cwd=project_directory,
keep_alive=keep_alive,
)
self.tool_name = tool_name
@ -419,6 +514,7 @@ class NpxStdioTransport(StdioTransport):
project_directory: str | None = None,
env_vars: dict[str, str] | None = None,
use_package_lock: bool = True,
keep_alive: bool | None = None,
):
"""
Initialize an Npx transport.
@ -456,7 +552,13 @@ class NpxStdioTransport(StdioTransport):
env = os.environ.copy()
env.update(env_vars)
super().__init__(command="npx", args=npx_args, env=env, cwd=project_directory)
super().__init__(
command="npx",
args=npx_args,
env=env,
cwd=project_directory,
keep_alive=keep_alive,
)
self.package = package

129
tests/client/test_stdio.py Normal file
View file

@ -0,0 +1,129 @@
import inspect
import pytest
from mcp.types import TextContent
from fastmcp import Client
from fastmcp.client.transports import PythonStdioTransport, StdioTransport
class TestKeepAlive:
# https://github.com/jlowin/fastmcp/issues/581
@pytest.fixture
def stdio_script(self, tmp_path):
script = inspect.cleandoc('''
import os
from fastmcp import FastMCP
mcp = FastMCP()
@mcp.tool()
def pid() -> int:
"""Gets PID of server"""
return os.getpid()
if __name__ == "__main__":
mcp.run()
''')
script_file = tmp_path / "stdio.py"
script_file.write_text(script)
return script_file
async def test_keep_alive_default_true(self):
client = Client(transport=StdioTransport(command="python", args=[""]))
assert client.transport.keep_alive is True
async def test_keep_alive_set_false(self):
client = Client(
transport=StdioTransport(command="python", args=[""], keep_alive=False)
)
assert client.transport.keep_alive is False
async def test_keep_alive_maintains_session_across_multiple_calls(
self, stdio_script
):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
assert isinstance(result1[0], TextContent)
pid1 = int(result1[0].text)
async with client:
result2 = await client.call_tool("pid")
assert isinstance(result2[0], TextContent)
pid2 = int(result2[0].text)
assert pid1 == pid2
async def test_keep_alive_false_starts_new_session_across_multiple_calls(
self, stdio_script
):
client = Client(
transport=PythonStdioTransport(script_path=stdio_script, keep_alive=False)
)
assert client.transport.keep_alive is False
async with client:
result1 = await client.call_tool("pid")
assert isinstance(result1[0], TextContent)
pid1 = int(result1[0].text)
async with client:
result2 = await client.call_tool("pid")
assert isinstance(result2[0], TextContent)
pid2 = int(result2[0].text)
assert pid1 != pid2
async def test_keep_alive_starts_new_session_if_manually_closed(self, stdio_script):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
assert isinstance(result1[0], TextContent)
pid1 = int(result1[0].text)
await client.close()
async with client:
result2 = await client.call_tool("pid")
assert isinstance(result2[0], TextContent)
pid2 = int(result2[0].text)
assert pid1 != pid2
async def test_keep_alive_maintains_session_if_reentered(self, stdio_script):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
async with client:
result1 = await client.call_tool("pid")
assert isinstance(result1[0], TextContent)
pid1 = int(result1[0].text)
async with client:
result2 = await client.call_tool("pid")
assert isinstance(result2[0], TextContent)
pid2 = int(result2[0].text)
result3 = await client.call_tool("pid")
assert isinstance(result3[0], TextContent)
pid3 = int(result3[0].text)
assert pid1 == pid2 == pid3
async def test_close_session_and_try_to_use_client_raises_error(self, stdio_script):
client = Client(transport=PythonStdioTransport(script_path=stdio_script))
assert client.transport.keep_alive is True
with pytest.raises(
RuntimeError, match="Server session was closed unexpectedly"
):
async with client:
await client.close()
await client.call_tool("pid")