diff --git a/tests/tools/test_mcp_multiplex_connection_keys.py b/tests/tools/test_mcp_multiplex_connection_keys.py index adaa1ca1b0..faaa0aad7e 100644 --- a/tests/tools/test_mcp_multiplex_connection_keys.py +++ b/tests/tools/test_mcp_multiplex_connection_keys.py @@ -307,12 +307,8 @@ def test_launch_profile_pruning_a_server_keeps_served_profiles_same_named_connec ``reconcile_mcp_servers_with_config`` prunes with ``shutdown_mcp_servers(scope=None, names={"x"})``. ``scope=None`` must mean *the unscoped owner* there, not *every owner* — otherwise the dashboard's own profile silently tears down profile B's ``(B, "x")``.""" - import asyncio - import threading - import tools.mcp_tool as core - from tools import mcp_tool_discovery as disc, mcp_tool_lifecycle as lifecycle - from tools import mcp_tool_registration as reg + from tools import mcp_tool_discovery as disc, mcp_tool_lifecycle as lifecycle, mcp_tool_loop as loop monkeypatch.setattr("agent.secret_scope.is_multiplex_active", lambda: False) cfg = {"url": "https://mcp.example/x", "headers": {"Authorization": "Bearer shared"}} @@ -320,7 +316,6 @@ def test_launch_profile_pruning_a_server_keeps_served_profiles_same_named_connec scope_b = two_profiles("b") srv_b = _server("x", cfg) disc._adopt_server("x", srv_b) - srv_b._registered_tool_names = reg._register_server_tools("x", srv_b, cfg) assert core._server_scope_keys[(scope_b, "x")] == scope_b with patch("hermes_constants.get_hermes_home_override", return_value=None): @@ -328,6 +323,9 @@ def test_launch_profile_pruning_a_server_keeps_served_profiles_same_named_connec srv_launch = _server("x", cfg) disc._adopt_server("x", srv_launch) assert core._server_scope_keys["x"] is None + # A third profile adopted the launch profile's connection: pruning it must remember the + # adopter so the next discovery pass re-registers it. + core._server_tool_scopes["x"] = {"c"} closed = [] @@ -337,19 +335,14 @@ def test_launch_profile_pruning_a_server_keeps_served_profiles_same_named_connec for srv in (srv_b, srv_launch): srv.shutdown = _shutdown.__get__(srv) - loop = asyncio.new_event_loop() - thread = threading.Thread(target=loop.run_forever, daemon=True) - thread.start() - monkeypatch.setattr(core, "_mcp_loop", loop) + loop._ensure_mcp_loop() try: - with patch.object(lifecycle._loop, "_stop_mcp_loop", lambda **_kw: False): - lifecycle.shutdown_mcp_servers(scope=None, names={"x"}) + lifecycle.shutdown_mcp_servers(scope=None, names={"x"}) finally: - loop.call_soon_threadsafe(loop.stop) - thread.join(timeout=5) - loop.close() + loop._stop_mcp_loop() assert "x" not in core._servers and "x" not in core._server_scope_keys assert core._servers[(scope_b, "x")] is srv_b assert core._server_scope_keys[(scope_b, "x")] == scope_b assert closed == ["x"] + assert core._orphaned_adopters == {"c": {"x"}} diff --git a/tools/mcp_tool_lifecycle.py b/tools/mcp_tool_lifecycle.py index e986fbd634..3b9fe5f07e 100644 --- a/tools/mcp_tool_lifecycle.py +++ b/tools/mcp_tool_lifecycle.py @@ -146,7 +146,7 @@ def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] = servers_snapshot = [_core._servers[key] for key in selected] if names is not None: selected_status = set(selected) - elif scope is None: + elif wildcard: selected_status = ( set(_core._servers) | set(_core._server_scope_keys) | set(_core._server_tool_scopes) @@ -184,7 +184,7 @@ def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] = _core._servers.pop(key, None) _core._server_scope_keys.pop(key, None) clear_selected_status() - _clear_connect_cooldowns(None if scope is None and names is None else selected_status) + _clear_connect_cooldowns(None if wildcard else selected_status) with _core._lock: loop = _core._mcp_loop @@ -203,8 +203,8 @@ def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] = with _core._lock: if not servers_snapshot: clear_selected_status() - _clear_connect_cooldowns(None if scope is None and names is None else selected_status) - _loop._stop_mcp_loop(only_if_idle=scope is not None or names is not None) + _clear_connect_cooldowns(None if wildcard else selected_status) + _loop._stop_mcp_loop(only_if_idle=not wildcard) # A removed subset still shares its profile's log with the remaining servers. # Full/profile shutdown must also release handles left by completed CLI/UI probes. if names is None: