fix(mcp): clear connect-cooldown state on every shutdown path

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.
This commit is contained in:
Teknium 2026-07-21 05:42:48 -07:00
parent 2c3aa3f7b8
commit ee23bffee9
2 changed files with 57 additions and 1 deletions

View file

@ -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 == {}

View file

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