diff --git a/fastmcp_slim/fastmcp/client/transports/stdio.py b/fastmcp_slim/fastmcp/client/transports/stdio.py index a0067bfa0..4df89fc6f 100644 --- a/fastmcp_slim/fastmcp/client/transports/stdio.py +++ b/fastmcp_slim/fastmcp/client/transports/stdio.py @@ -83,6 +83,14 @@ class StdioTransport(ClientTransport): async def connect( self, **session_kwargs: Unpack[SessionKwargs] ) -> ClientSession | None: + current_loop = asyncio.get_running_loop() + if ( + self._connect_task is not None + and self._connect_task.get_loop() is not current_loop + ): + logger.debug("Stdio transport changed event loops; reconnecting") + await self.disconnect() + # If the connect task completed or the session's streams are dead, # the subprocess has exited. Tear down so we can start fresh. if self._connect_task is not None and ( @@ -127,6 +135,19 @@ class StdioTransport(ClientTransport): if self._connect_task is None: return + owner_loop = self._connect_task.get_loop() + current_loop = asyncio.get_running_loop() + if ( + owner_loop is not current_loop + and not self._connect_task.done() + and owner_loop.is_running() + ): + disconnect_future = asyncio.run_coroutine_threadsafe( + self.disconnect(), owner_loop + ) + await asyncio.wrap_future(disconnect_future) + return + # signal the connection task to stop self._stop_event.set() diff --git a/tests/client/test_stdio.py b/tests/client/test_stdio.py index af4995203..dfac9fc23 100644 --- a/tests/client/test_stdio.py +++ b/tests/client/test_stdio.py @@ -2,6 +2,7 @@ import asyncio import gc import inspect import os +import threading import weakref import psutil @@ -237,6 +238,44 @@ class TestKeepAlive: assert pid1 == pid2 == pid3 + def test_keep_alive_reconnects_on_different_running_loop(self, stdio_script): + transport = PythonStdioTransport(script_path=stdio_script) + client = Client(transport=transport, init_timeout=2) + first_loop = asyncio.new_event_loop() + ready = threading.Event() + + def run_first_loop(): + asyncio.set_event_loop(first_loop) + ready.set() + first_loop.run_forever() + + async def get_pid(*, close: bool = False) -> int: + async with client: + result = await client.call_tool("pid") + if close: + await client.close() + return result.data + + thread = threading.Thread(target=run_first_loop) + thread.start() + try: + assert ready.wait(timeout=2) + pid1 = asyncio.run_coroutine_threadsafe(get_pid(), first_loop).result( + timeout=5 + ) + pid2 = asyncio.run(asyncio.wait_for(get_pid(close=True), timeout=5)) + + assert pid1 != pid2 + finally: + if transport._connect_task is not None: + owner_loop = transport._connect_task.get_loop() + asyncio.run_coroutine_threadsafe(transport.close(), owner_loop).result( + timeout=5 + ) + first_loop.call_soon_threadsafe(first_loop.stop) + thread.join(timeout=2) + first_loop.close() + 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