Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 5 additions & 6 deletions cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
12 changes: 5 additions & 7 deletions gateway/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
9 changes: 3 additions & 6 deletions tests/gateway/test_mcp_reload_refreshes_cached_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
):
Expand Down Expand Up @@ -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=[]),
):
Expand All @@ -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),
):
Expand Down
64 changes: 63 additions & 1 deletion tests/tools/test_mcp_stability.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
123 changes: 92 additions & 31 deletions tools/mcp_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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":
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
5 changes: 2 additions & 3 deletions tui_gateway/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading