From 4f8309e254f9ce8904111ed3518f696d9692a83f Mon Sep 17 00:00:00 2001
From: Flownium <157689911+itsflownium@users.noreply.github.com>
Date: Mon, 15 Jun 2026 17:11:01 +1000
Subject: [PATCH] fix: restart mcp loop during reload

---
 cli.py                                        |  11 +-
 gateway/run.py                                |  12 +-
 ...test_mcp_reload_refreshes_cached_agents.py |   9 +-
 tests/tools/test_mcp_stability.py             |  64 ++++++++-
 tools/mcp_tool.py                             | 123 +++++++++++++-----
 tui_gateway/server.py                         |   5 +-
 6 files changed, 170 insertions(+), 54 deletions(-)

diff --git a/cli.py b/cli.py
index 25cce4f95d05f..6545b93a56eb2 100644
--- a/cli.py
+++ b/cli.py
@@ -10792,7 +10792,7 @@ def _reload_mcp(self):
         sees the updated tools on the next turn.
         """
         try:
-            from tools.mcp_tool import shutdown_mcp_servers, discover_mcp_tools, _servers, _lock
+            from tools.mcp_tool import reload_mcp_servers, _servers, _lock
 
             # Capture old server names
             with _lock:
@@ -10801,11 +10801,10 @@ def _reload_mcp(self):
             if not self._command_running:
                 print("🔄 Reloading MCP servers...")
 
-            # Shutdown existing connections
-            shutdown_mcp_servers()
-
-            # Reconnect (reads config.yaml fresh)
-            new_tools = discover_mcp_tools()
+            # Shutdown and reconnect (reads config.yaml fresh) as one
+            # serialized operation so stdio transports finish unwinding
+            # cancellation before discovery schedules fresh server tasks.
+            new_tools = reload_mcp_servers()
 
             # Compute what changed
             with _lock:
diff --git a/gateway/run.py b/gateway/run.py
index f584089ebffde..4e94e54404762 100644
--- a/gateway/run.py
+++ b/gateway/run.py
@@ -14203,18 +14203,16 @@ async def _execute_mcp_reload(self, event: MessageEvent) -> str:
         """
         loop = asyncio.get_running_loop()
         try:
-            from tools.mcp_tool import shutdown_mcp_servers, discover_mcp_tools, _servers, _lock
+            from tools.mcp_tool import reload_mcp_servers, _servers, _lock
 
             # Capture old server names before shutdown
             with _lock:
                 old_servers = set(_servers.keys())
 
-            # Read new config before shutting down, so we know what will be added/removed
-            # Shutdown existing connections
-            await loop.run_in_executor(None, shutdown_mcp_servers)
-
-            # Reconnect by discovering tools (reads config.yaml fresh)
-            new_tools = await loop.run_in_executor(None, discover_mcp_tools)
+            # Shutdown and reconnect (reads config.yaml fresh) as one
+            # serialized operation so stdio transports finish unwinding
+            # cancellation before discovery schedules fresh server tasks.
+            new_tools = await loop.run_in_executor(None, reload_mcp_servers)
 
             # Compute what changed
             with _lock:
diff --git a/tests/gateway/test_mcp_reload_refreshes_cached_agents.py b/tests/gateway/test_mcp_reload_refreshes_cached_agents.py
index 4d945f03c5942..6f887276e4bc8 100644
--- a/tests/gateway/test_mcp_reload_refreshes_cached_agents.py
+++ b/tests/gateway/test_mcp_reload_refreshes_cached_agents.py
@@ -106,8 +106,7 @@ async def test_reload_mcp_refreshes_cached_agent_tools():
     ]
 
     with (
-        patch("tools.mcp_tool.shutdown_mcp_servers"),
-        patch("tools.mcp_tool.discover_mcp_tools", return_value=["HassTurnOn", "HassTurnOff"]),
+        patch("tools.mcp_tool.reload_mcp_servers", return_value=["HassTurnOn", "HassTurnOff"]),
         patch.dict("tools.mcp_tool._servers", {"homeassistant": object()}, clear=True),
         patch("model_tools.get_tool_definitions", return_value=fresh_tool_defs),
     ):
@@ -136,8 +135,7 @@ async def test_reload_mcp_handles_empty_agent_cache():
     assert len(runner._agent_cache) == 0
 
     with (
-        patch("tools.mcp_tool.shutdown_mcp_servers"),
-        patch("tools.mcp_tool.discover_mcp_tools", return_value=[]),
+        patch("tools.mcp_tool.reload_mcp_servers", return_value=[]),
         patch.dict("tools.mcp_tool._servers", {}, clear=True),
         patch("model_tools.get_tool_definitions", return_value=[]),
     ):
@@ -164,8 +162,7 @@ def _capture_get_tool_definitions(**kwargs):
         return [{"type": "function", "function": {"name": "refreshed"}}]
 
     with (
-        patch("tools.mcp_tool.shutdown_mcp_servers"),
-        patch("tools.mcp_tool.discover_mcp_tools", return_value=["refreshed"]),
+        patch("tools.mcp_tool.reload_mcp_servers", return_value=["refreshed"]),
         patch.dict("tools.mcp_tool._servers", {"homeassistant": object()}, clear=True),
         patch("model_tools.get_tool_definitions", side_effect=_capture_get_tool_definitions),
     ):
diff --git a/tests/tools/test_mcp_stability.py b/tests/tools/test_mcp_stability.py
index b9a1b92c9cdea..7ff81d66f88d4 100644
--- a/tests/tools/test_mcp_stability.py
+++ b/tests/tools/test_mcp_stability.py
@@ -3,7 +3,7 @@
 import asyncio
 import os
 import signal
-from unittest.mock import patch, MagicMock
+from unittest.mock import AsyncMock, patch, MagicMock
 
 import pytest
 
@@ -57,6 +57,68 @@ def test_handler_installed_on_mcp_loop(self):
         finally:
             mcp_mod._stop_mcp_loop()
 
+    def test_reload_mcp_servers_discovers_on_fresh_loop(self):
+        """Full reload stops the old MCP loop before scheduling discovery."""
+        import tools.mcp_tool as mcp_mod
+
+        with mcp_mod._lock:
+            mcp_mod._servers.clear()
+            mcp_mod._server_connecting.clear()
+            mcp_mod._server_connect_errors.clear()
+
+        mcp_mod._ensure_mcp_loop()
+        with mcp_mod._lock:
+            old_loop = mcp_mod._mcp_loop
+            old_thread = mcp_mod._mcp_thread
+            old_server = MagicMock()
+            old_server.name = "old"
+            old_server.shutdown = AsyncMock()
+            mcp_mod._servers["old"] = old_server
+
+        discovery_loops = []
+
+        async def _fake_discover(name, cfg):
+            discovery_loops.append(asyncio.get_running_loop())
+            server = MagicMock()
+            server.name = name
+            server._registered_tool_names = ["mcp_fresh_ping"]
+            with mcp_mod._lock:
+                mcp_mod._servers[name] = server
+            return ["mcp_fresh_ping"]
+
+        try:
+            with patch("tools.mcp_tool._MCP_AVAILABLE", True), \
+                 patch("tools.mcp_tool._load_mcp_config", return_value={"fresh": {"command": "python3"}}), \
+                 patch("tools.mcp_tool._discover_and_register_server", side_effect=_fake_discover), \
+                 patch("tools.mcp_tool._existing_tool_names", return_value=["mcp_fresh_ping"]):
+                result = mcp_mod.reload_mcp_servers()
+
+            assert result == ["mcp_fresh_ping"]
+            old_server.shutdown.assert_awaited_once()
+            assert discovery_loops
+            assert discovery_loops[0] is not old_loop
+            assert old_thread is not None
+            assert not old_thread.is_alive()
+            with mcp_mod._lock:
+                assert set(mcp_mod._servers) == {"fresh"}
+        finally:
+            with mcp_mod._lock:
+                mcp_mod._servers.clear()
+                mcp_mod._server_connecting.clear()
+                mcp_mod._server_connect_errors.clear()
+            mcp_mod._stop_mcp_loop()
+
+    def test_reload_mcp_servers_refuses_discovery_when_shutdown_fails(self):
+        """A partially stopped loop must not be reused for discovery."""
+        import tools.mcp_tool as mcp_mod
+
+        with patch.object(mcp_mod, "shutdown_mcp_servers", return_value=False), \
+             patch.object(mcp_mod, "discover_mcp_tools") as discover:
+            with pytest.raises(RuntimeError, match="did not stop cleanly"):
+                mcp_mod.reload_mcp_servers()
+
+        discover.assert_not_called()
+
     def test_probe_cleanup_does_not_stop_loop_with_registered_servers(self):
         """Probe cleanup must not kill the shared loop used by live MCP tools."""
         import tools.mcp_tool as mcp_mod
diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py
index 4438a715718a1..3bc1e232f87ee 100644
--- a/tools/mcp_tool.py
+++ b/tools/mcp_tool.py
@@ -3604,11 +3604,18 @@ def _handle_session_expired_and_retry(
 # Dedicated event loop running in a background daemon thread.
 _mcp_loop: Optional[asyncio.AbstractEventLoop] = None
 _mcp_thread: Optional[threading.Thread] = None
+_MCP_LOOP_START_TIMEOUT = 2.0
+_MCP_LOOP_STOP_TIMEOUT = 15.0
 
 # Protects _mcp_loop, _mcp_thread, _servers, MCP connection status maps,
 # _parallel_safe_servers, _mcp_tool_server_names, and _stdio_pids.
 _lock = threading.Lock()
 
+# Serializes loop start/stop/reload so a rediscovery cannot schedule work onto
+# a loop while the previous stdio transports are still unwinding cancellation.
+_mcp_lifecycle_lock = threading.RLock()
+_mcp_reload_lock = threading.Lock()
+
 # PIDs of stdio MCP server subprocesses.  Tracked so we can force-kill
 # them on shutdown if the graceful cleanup (SDK context-manager teardown)
 # fails or times out.  PIDs are added after connection and removed on
@@ -3731,17 +3738,34 @@ def _mcp_loop_exception_handler(loop, context):
 def _ensure_mcp_loop():
     """Start the background event loop thread if not already running."""
     global _mcp_loop, _mcp_thread
-    with _lock:
-        if _mcp_loop is not None and _mcp_loop.is_running():
-            return
-        _mcp_loop = asyncio.new_event_loop()
-        _mcp_loop.set_exception_handler(_mcp_loop_exception_handler)
-        _mcp_thread = threading.Thread(
-            target=_mcp_loop.run_forever,
-            name="mcp-event-loop",
-            daemon=True,
-        )
-        _mcp_thread.start()
+    with _mcp_lifecycle_lock:
+        ready = threading.Event()
+
+        with _lock:
+            if _mcp_loop is not None and _mcp_loop.is_running():
+                return
+            loop = asyncio.new_event_loop()
+            loop.set_exception_handler(_mcp_loop_exception_handler)
+
+            def _run_loop() -> None:
+                asyncio.set_event_loop(loop)
+                loop.call_soon(ready.set)
+                loop.run_forever()
+
+            thread = threading.Thread(
+                target=_run_loop,
+                name="mcp-event-loop",
+                daemon=True,
+            )
+            _mcp_loop = loop
+            _mcp_thread = thread
+            thread.start()
+
+        if not ready.wait(timeout=_MCP_LOOP_START_TIMEOUT):
+            logger.warning(
+                "MCP event loop did not report ready within %.1fs",
+                _MCP_LOOP_START_TIMEOUT,
+            )
 
 
 def _wrap_with_home_override(coro: "Coroutine") -> "Coroutine":
@@ -5621,20 +5645,22 @@ def _add(schema: dict) -> bool:
     return staged_engine_names
 
 
-def shutdown_mcp_servers():
+def shutdown_mcp_servers() -> bool:
     """Close all MCP server connections and stop the background loop.
 
     Each server Task is signalled to exit its ``async with`` block so that
     the anyio cancel-scope cleanup happens in the same Task that opened it.
     All servers are shut down in parallel via ``asyncio.gather``.
+
+    Returns True when the MCP loop was stopped or was already absent. A False
+    return means the old loop thread did not finish within the bounded wait.
     """
     with _lock:
         servers_snapshot = list(_servers.values())
 
     # Fast path: nothing to shut down.
     if not servers_snapshot:
-        _stop_mcp_loop()
-        return
+        return _stop_mcp_loop()
 
     async def _shutdown():
         results = await asyncio.gather(
@@ -5664,7 +5690,25 @@ async def _shutdown():
             except BaseException as exc:
                 logger.debug("Error during MCP shutdown: %s", exc)
 
-    _stop_mcp_loop()
+    return _stop_mcp_loop()
+
+
+def reload_mcp_servers() -> List[str]:
+    """Synchronously shut down MCP servers and rediscover them on a fresh loop.
+
+    Reload callers must not stitch together ``shutdown_mcp_servers()`` and
+    ``discover_mcp_tools()`` themselves: stdio transports rely on anyio
+    cancellation scopes that need to finish unwinding before new server tasks
+    are scheduled.  Serializing the full reload here gives every UI entry point
+    the same clean shutdown -> fresh discovery boundary.
+    """
+    with _mcp_reload_lock:
+        if not shutdown_mcp_servers():
+            raise RuntimeError(
+                "MCP event loop did not stop cleanly; refusing to rediscover "
+                "servers on a partially shut down loop"
+            )
+        return discover_mcp_tools()
 
 
 def _kill_orphaned_mcp_children(
@@ -5810,24 +5854,41 @@ def _stop_mcp_loop_if_idle() -> bool:
 def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool:
     """Stop the background event loop and join its thread."""
     global _mcp_loop, _mcp_thread
-    with _lock:
-        if only_if_idle and (_servers or _server_connecting):
-            logger.debug("Leaving MCP event loop running; active servers are registered or connecting")
-            return False
-        loop = _mcp_loop
-        thread = _mcp_thread
-        _mcp_loop = None
-        _mcp_thread = None
-    if loop is not None:
-        loop.call_soon_threadsafe(loop.stop)
-        if thread is not None:
-            thread.join(timeout=5)
+    with _mcp_lifecycle_lock:
+        with _lock:
+            if only_if_idle and (_servers or _server_connecting):
+                logger.debug("Leaving MCP event loop running; active servers are registered or connecting")
+                return False
+            loop = _mcp_loop
+            thread = _mcp_thread
+            _mcp_loop = None
+            _mcp_thread = None
+
+        if loop is None:
+            return True
+
         try:
-            loop.close()
-        except Exception:
-            pass
+            if loop.is_running():
+                loop.call_soon_threadsafe(loop.stop)
+        except RuntimeError as exc:
+            logger.debug("MCP event loop stop scheduling failed: %s", exc)
+
+        if thread is not None and thread is not threading.current_thread():
+            thread.join(timeout=_MCP_LOOP_STOP_TIMEOUT)
+            if thread.is_alive():
+                logger.warning(
+                    "MCP event loop thread did not stop within %.1fs",
+                    _MCP_LOOP_STOP_TIMEOUT,
+                )
+                return False
+
+        if not loop.is_closed():
+            try:
+                loop.close()
+            except Exception:
+                pass
         # After closing the loop, any stdio subprocesses that survived the
         # graceful shutdown are now orphaned — include active PIDs too
         # since the loop is gone and no session can still be in flight.
         _kill_orphaned_mcp_children(include_active=True)
-    return True
+        return True
diff --git a/tui_gateway/server.py b/tui_gateway/server.py
index af5cead103c1a..744ea36302c30 100644
--- a/tui_gateway/server.py
+++ b/tui_gateway/server.py
@@ -11622,10 +11622,9 @@ def _(rid, params: dict) -> dict:
                     },
                 )
 
-        from tools.mcp_tool import shutdown_mcp_servers, discover_mcp_tools
+        from tools.mcp_tool import reload_mcp_servers
 
-        shutdown_mcp_servers()
-        discover_mcp_tools()
+        reload_mcp_servers()
         if session:
             agent = session["agent"]
             # Rebuild the cached agent's tool snapshot so the current session