Reconnect stdio transport across event loops

🤖 Generated with Codex

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

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