From ee23bffee9acf31b61c32f63d01c7996ea711c72 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Tue, 21 Jul 2026 05:42:48 -0700
Subject: [PATCH] fix(mcp): clear connect-cooldown state on every shutdown path
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Builds on trevorgordon981's #50589 (cherry-picked as the previous commit).
The #50394 cooldown reset only ran inside the async _shutdown coroutine,
which is skipped on the empty-_servers fast path — the most common state
when a server failed to connect (failed servers are never recorded in
_servers). It was also skipped when the MCP loop wasn't running.

Clear _server_connect_retry_after/_server_connect_failures on the fast
path and in a final unconditional sweep so a full shutdown/restart always
re-attempts every configured server immediately. Adds regression tests
for both paths.
---
 tests/tools/test_mcp_bridge_single_failure.py | 40 +++++++++++++++++++
 tools/mcp_tool.py                             | 18 ++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/tests/tools/test_mcp_bridge_single_failure.py b/tests/tools/test_mcp_bridge_single_failure.py
index 3be44d41685..2185510d874 100644
--- a/tests/tools/test_mcp_bridge_single_failure.py
+++ b/tests/tools/test_mcp_bridge_single_failure.py
@@ -130,3 +130,43 @@ class TestRegisterMcpServersIsolation:
             attempts.clear()
             mcp_mod.register_mcp_servers(cfg)
             assert "bad" in attempts, "elapsed cooldown should permit a retry"
+
+
+class TestShutdownClearsCooldownState:
+    """shutdown_mcp_servers must drop cooldown state on EVERY path.
+
+    A server that failed to connect is never recorded in ``_servers``, so
+    the empty-``_servers`` fast path is the most common state in which
+    stale cooldown entries exist. The original #50394 fix only cleared the
+    maps inside the async ``_shutdown`` coroutine, which the fast path
+    (and the loop-not-running path) never executes.
+    """
+
+    def test_fast_path_clears_cooldown_state(self):
+        mcp_mod._record_connect_failure("bad")
+        assert mcp_mod._server_connect_retry_after
+        assert not mcp_mod._servers  # precondition: fast path taken
+
+        with patch("tools.mcp_tool._stop_mcp_loop"):
+            mcp_mod.shutdown_mcp_servers()
+
+        assert mcp_mod._server_connect_retry_after == {}
+        assert mcp_mod._server_connect_failures == {}
+
+    def test_loop_not_running_path_clears_cooldown_state(self):
+        mcp_mod._record_connect_failure("bad")
+
+        class _DeadServer:
+            name = "dead"
+
+            async def shutdown(self):  # pragma: no cover - never awaited
+                pass
+
+        mcp_mod._servers["dead"] = _DeadServer()  # type: ignore[assignment]
+        # _mcp_loop is None in this test process, so the async _shutdown
+        # coroutine is never scheduled; only the final sweep can clear.
+        with patch("tools.mcp_tool._stop_mcp_loop"):
+            mcp_mod.shutdown_mcp_servers()
+
+        assert mcp_mod._server_connect_retry_after == {}
+        assert mcp_mod._server_connect_failures == {}
diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py
index af14b122a1d..1097968df4a 100644
--- a/tools/mcp_tool.py
+++ b/tools/mcp_tool.py
@@ -6102,8 +6102,16 @@ def shutdown_mcp_servers():
     with _lock:
         servers_snapshot = list(_servers.values())
 
-    # Fast path: nothing to shut down.
+    # Fast path: nothing to shut down. The connect-cooldown maps can still
+    # be populated here — a server that failed to connect is never recorded
+    # in ``_servers`` (that is the very premise of the #50394 cooldown), so
+    # "no live servers" is the MOST likely state in which stale backoff
+    # entries exist. Clear them so a post-shutdown restart re-attempts every
+    # configured server immediately.
     if not servers_snapshot:
+        with _lock:
+            _server_connect_retry_after.clear()
+            _server_connect_failures.clear()
         _stop_mcp_loop()
         return
 
@@ -6140,6 +6148,14 @@ def shutdown_mcp_servers():
             except BaseException as exc:
                 logger.debug("Error during MCP shutdown: %s", exc)
 
+    # Unconditional final sweep: whether the async ``_shutdown`` ran,
+    # timed out, or was never scheduled (loop already stopped), a full
+    # shutdown must leave no stale connect-cooldown state behind — the
+    # next start should re-attempt every server immediately (#50394).
+    with _lock:
+        _server_connect_retry_after.clear()
+        _server_connect_failures.clear()
+
     _stop_mcp_loop()