fix(mcp): divide the shutdown budget across profile passes and stop poisoning the teardown context

Each per-profile pass could spend 15s against a 5s caller budget, so on a host with
several served profiles the trailing wildcard pass — the only one that stops the shared
MCP loop and its connections — might never run. shutdown_mcp_servers takes a timeout and
the caller gives each pass budget/(N+1).

The worker also ran under copy_context(), carrying a caller's HERMES_HOME override into
a pass that documents it has none: from inside profiles/poison's scope the wildcard pass
resolved live_home to that profile. It now starts from a fresh Context().

Both call sites replace `with suppress(Exception)` with a logged WARNING: the suppress
swallowed a real TypeError from this signature change, which surfaced only as a missing
teardown in an unrelated test.
This commit is contained in:
teknium1
2026-09-21 02:50:34 -07:00
committed by Teknium
parent bc0a42fd96
commit af380f66ae
3 changed files with 30 additions and 11 deletions

View File

@@ -24,7 +24,7 @@ import threading
import time
import traceback
from collections import OrderedDict
from contextvars import copy_context
from contextvars import Context, copy_context
from pathlib import Path
from datetime import datetime
from typing import Callable, Dict, Optional, Any, List, Tuple, cast
@@ -4847,12 +4847,21 @@ async def _shutdown_mcp_servers_nonblocking(timeout: float = 5.0, config: Any =
wildcard call (under the launch profile's own scope) stops the shared loop and reaps anything
the per-profile passes did not own.
``timeout`` is a TOTAL budget: each pass gets ``timeout / (N + 1)``, because the default 15s
per-pass wait inside ``shutdown_mcp_servers`` let N profiles consume the whole caller budget and
the trailing wildcard pass — the only one that stops the shared loop — never ran.
The worker runs in a FRESH context, not ``copy_context()``: the caller may sit inside a served
profile's scope, and ``launch_profile_scope_if_multiplexed`` documents "no HERMES_HOME override"
— inheriting one made the wildcard pass resolve the live home to that profile.
See #82874.
"""
from tui_gateway.launch_profile_policy import launch_profile_scope_if_multiplexed
profile_homes = (
_multiplex_profile_homes(config) if getattr(config, "multiplex_profiles", False) else [])
pass_timeout = max(1.0, timeout / (len(profile_homes) + 1))
def _do() -> None:
from tools.mcp_tool_common import _core
@@ -4860,16 +4869,16 @@ async def _shutdown_mcp_servers_nonblocking(timeout: float = 5.0, config: Any =
for profile_name, profile_home in profile_homes:
try:
with _profile_runtime_scope(Path(profile_home), hydrate_secrets=False):
shutdown_mcp_servers(scope=_core._mcp_registry_scope())
shutdown_mcp_servers(scope=_core._mcp_registry_scope(), timeout=pass_timeout)
except Exception:
logger.debug("MCP shutdown raised for profile '%s'", profile_name, exc_info=True)
try:
with launch_profile_scope_if_multiplexed():
shutdown_mcp_servers()
shutdown_mcp_servers(timeout=pass_timeout)
except Exception:
logger.debug("MCP shutdown raised", exc_info=True)
thread = threading.Thread(target=copy_context().run, args=(_do,), name="mcp-shutdown", daemon=True)
thread = threading.Thread(target=Context().run, args=(_do,), name="mcp-shutdown", daemon=True)
thread.start()
done = await _await_thread_exit(thread, timeout=timeout)
if not done:
@@ -5474,8 +5483,12 @@ async def _start_gateway_shutdown_tail(
_planned_stop_watcher_stop.set()
_planned_stop_watcher_thread.join(timeout=2)
with suppress(Exception):
# Never suppressed: a raise here is a real teardown failure (it once hid a changed signature,
# leaving every MCP connection and the shared loop up while the gateway reported a clean exit).
try:
await _shutdown_mcp_servers_nonblocking(config=getattr(runner, "config", None))
except Exception:
logger.warning("MCP shutdown failed; connections may be left open", exc_info=True)
# The failure verdict comes AFTER the cooperative teardown: returning early here leaked the
# cron ticker + housekeeping threads (and open MCP connections) for embedded callers (#12175).
@@ -5621,8 +5634,10 @@ async def start_gateway(config: Optional[GatewayConfig] = None, replace: bool =
# Startup aborted by restart/shutdown before running mode; preserve that path without starting cron.
try:
await runner.wait_for_shutdown()
with suppress(Exception):
try:
await _shutdown_mcp_servers_nonblocking(config=getattr(runner, "config", None))
except Exception:
logger.warning("MCP shutdown failed; connections may be left open", exc_info=True)
return _resolve_gateway_exit_verdict(runner, _signal_initiated_shutdown[0])
finally:
_shutdown_gateway_health_export(runner)

View File

@@ -407,7 +407,7 @@ async def test_shutdown_mcp_servers_nonblocking_keeps_loop_responsive():
started = asyncio.Event()
loop = asyncio.get_running_loop()
def wedged_shutdown():
def wedged_shutdown(**_kwargs): # the caller divides its budget across per-profile passes
loop.call_soon_threadsafe(started.set)
import time as _time
@@ -440,7 +440,7 @@ async def test_shutdown_mcp_servers_nonblocking_keeps_loop_responsive():
@pytest.mark.asyncio
async def test_shutdown_mcp_servers_nonblocking_completes_fast_path():
calls = []
with patch("tools.mcp_tool_lifecycle.shutdown_mcp_servers", lambda: calls.append(1)):
with patch("tools.mcp_tool_lifecycle.shutdown_mcp_servers", lambda **_kw: calls.append(1)):
done = await gateway_run._shutdown_mcp_servers_nonblocking(timeout=5)
assert done is True
assert calls == [1]

View File

@@ -127,7 +127,8 @@ def _reregister_orphaned_adopters() -> None:
reset_hermes_home_override(home_token)
def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] = None):
def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] = None,
timeout: float = 15.0):
"""Close MCP server connections (in parallel) and stop the background loop. Each server
Task is signalled to exit its own ``async with`` so the anyio cancel-scope cleanup runs in
the Task that opened it. ``scope`` restricts teardown to one multiplexed profile's servers
@@ -136,7 +137,10 @@ def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] =
(dropped-from-config pruning); other servers' bookkeeping is untouched. Only the bare call
(no ``scope``, no ``names``) is the process-wide wildcard: the launch profile's registry
scope IS ``None``, so ``scope=None, names={...}`` prunes that unscoped owner's servers and
must leave a served profile's same-named ``(B, name)`` connection alone."""
must leave a served profile's same-named ``(B, name)`` connection alone. ``timeout`` bounds
the wait for the close to land on the MCP loop — a caller running one pass per served
profile under a total budget divides it, or N profiles × 15s starve the wildcard pass that
actually stops the loop."""
from tools.mcp_tool_scope import _key_name
wildcard = scope is None and names is None
with _core._lock:
@@ -193,7 +197,7 @@ def shutdown_mcp_servers(*, scope: Optional[str] = None, names: Optional[set] =
future = safe_schedule_threadsafe(_shutdown(), loop, logger=logger, log_message="MCP shutdown: failed to schedule")
if future is not None:
try:
future.result(timeout=15)
future.result(timeout=timeout)
except BaseException as exc:
logger.debug("Error during MCP shutdown: %s", exc)