From 9bbd4171295bdbeecf6ad351d35ba5bdfa26cc51 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Wed, 28 May 2025 21:46:40 -0400
Subject: [PATCH] Add keep_alive param to reuse subprocess

---
 docs/clients/client.mdx          |  70 +++++++++++++----
 docs/clients/transports.mdx      |  59 +++++++++++++-
 src/fastmcp/client/client.py     |   6 ++
 src/fastmcp/client/transports.py | 130 +++++++++++++++++++++++++++----
 tests/client/test_stdio.py       | 129 ++++++++++++++++++++++++++++++
 5 files changed, 366 insertions(+), 28 deletions(-)
 create mode 100644 tests/client/test_stdio.py

diff --git a/docs/clients/client.mdx b/docs/clients/client.mdx
index c53ecee03..74ee9e3b4 100644
--- a/docs/clients/client.mdx
+++ b/docs/clients/client.mdx
@@ -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" />
diff --git a/docs/clients/transports.mdx b/docs/clients/transports.mdx
index 3669c8b25..6ed9fc489 100644
--- a/docs/clients/transports.mdx
+++ b/docs/clients/transports.mdx
@@ -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)
 
diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py
index a33969c8a..1fee12ff7 100644
--- a/src/fastmcp/client/client.py
+++ b/src/fastmcp/client/client.py
@@ -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:
diff --git a/src/fastmcp/client/transports.py b/src/fastmcp/client/transports.py
index 9a7268340..1fa52bf1a 100644
--- a/src/fastmcp/client/transports.py
+++ b/src/fastmcp/client/transports.py
@@ -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
 
 
diff --git a/tests/client/test_stdio.py b/tests/client/test_stdio.py
new file mode 100644
index 000000000..55d430f7c
--- /dev/null
+++ b/tests/client/test_stdio.py
@@ -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")