mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-30 19:09:28 +00:00
Merge remote-tracking branch 'origin/main' into tests/prune-low-value
# Conflicts: # tests/agent/test_context_compressor.py # tests/gateway/test_startup_restart_race.py # tests/hermes_cli/test_voice_wrapper.py
This commit is contained in:
commit
a17ac2ca67
7 changed files with 750 additions and 15 deletions
1
contributors/emails/amdnative@gmail.com
Normal file
1
contributors/emails/amdnative@gmail.com
Normal file
|
|
@ -0,0 +1 @@
|
|||
adurham
|
||||
1
contributors/emails/mrz@mrzlab630.pw
Normal file
1
contributors/emails/mrz@mrzlab630.pw
Normal file
|
|
@ -0,0 +1 @@
|
|||
mrzlab630
|
||||
|
|
@ -47,6 +47,7 @@ LEGACY_AUTHOR_MAP = {
|
|||
"declanbatesmith@outlook.com": "cat-thats-fat", # PR #60489 (desktop: first-run remote connection option)
|
||||
"drbs2004@me.com": "cat-thats-fat", # PR #60489 (desktop: first-run remote connection option; historical merge email)
|
||||
"122438640+ragingbulld@users.noreply.github.com": "ragingbulld", # PR #65606 salvage (non-finite API wait deadlines; #65746)
|
||||
"shady2k@gmail.com": "shady2k", # PR #60104 salvage of #66143 (MCP loop-owned shutdown drain)
|
||||
"zzpigpinggai@users.noreply.github.com": "zzpigpinggai", # PR #66017 salvage of #63617 (OpenRouter explicit-provider picker visibility)
|
||||
"stellarisw@users.noreply.github.com": "StellarisW", # PR #66222 salvage (Discord WebSocket liveness + systemd watchdog; #26656 follow-up)
|
||||
"wx.xw@bytedance.com": "wxy-nlp", # PR #66222 salvage (systemd event-loop watchdog co-author)
|
||||
|
|
|
|||
273
tests/tools/test_mcp_initial_connect_shutdown.py
Normal file
273
tests/tools/test_mcp_initial_connect_shutdown.py
Normal file
|
|
@ -0,0 +1,273 @@
|
|||
"""Regression tests for initial MCP failure ownership and teardown."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _reset_mcp_state(mcp_tool) -> None:
|
||||
mcp_tool.shutdown_mcp_servers()
|
||||
with mcp_tool._lock:
|
||||
mcp_tool._servers.clear()
|
||||
mcp_tool._server_connecting.clear()
|
||||
mcp_tool._server_connect_errors.clear()
|
||||
|
||||
|
||||
def _cleanup_mcp_state(mcp_tool, extra_servers=()) -> None:
|
||||
with mcp_tool._lock:
|
||||
loop = mcp_tool._mcp_loop
|
||||
if loop is not None and loop.is_running():
|
||||
for server in extra_servers:
|
||||
task = getattr(server, "_task", None)
|
||||
if task is not None and not task.done():
|
||||
mcp_tool._run_on_mcp_loop(server.shutdown, timeout=5)
|
||||
mcp_tool.shutdown_mcp_servers()
|
||||
with mcp_tool._lock:
|
||||
mcp_tool._servers.clear()
|
||||
mcp_tool._server_connecting.clear()
|
||||
mcp_tool._server_connect_errors.clear()
|
||||
|
||||
|
||||
def test_initial_connect_failure_is_registry_owned_and_reaped(monkeypatch, tmp_path):
|
||||
"""Normal discovery must retain the parked task for clean shutdown."""
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
||||
|
||||
from tools import mcp_tool
|
||||
|
||||
_reset_mcp_state(mcp_tool)
|
||||
created = []
|
||||
|
||||
class _FailingServerTask(mcp_tool.MCPServerTask):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
created.append(self)
|
||||
|
||||
async def _run_stdio(self, config):
|
||||
raise ConnectionError("deterministic initial failure")
|
||||
|
||||
monkeypatch.setattr(mcp_tool, "MCPServerTask", _FailingServerTask)
|
||||
monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True)
|
||||
monkeypatch.setattr(mcp_tool, "_MAX_INITIAL_CONNECT_RETRIES", 0)
|
||||
monkeypatch.setattr(mcp_tool, "_PARKED_RETRY_INTERVAL", 3600)
|
||||
|
||||
real_stop = mcp_tool._stop_mcp_loop
|
||||
pending_at_stop = []
|
||||
|
||||
async def _pending_tasks():
|
||||
current = asyncio.current_task()
|
||||
return sorted(
|
||||
task.get_coro().__qualname__
|
||||
for task in asyncio.all_tasks()
|
||||
if task is not current and not task.done()
|
||||
)
|
||||
|
||||
def _observed_stop(*, only_if_idle=False):
|
||||
pending_at_stop.extend(
|
||||
mcp_tool._run_on_mcp_loop(_pending_tasks, timeout=5)
|
||||
)
|
||||
return real_stop(only_if_idle=only_if_idle)
|
||||
|
||||
monkeypatch.setattr(mcp_tool, "_stop_mcp_loop", _observed_stop)
|
||||
|
||||
try:
|
||||
assert mcp_tool.register_mcp_servers({
|
||||
"initial-failure": {"command": "unused", "connect_timeout": 5}
|
||||
}) == []
|
||||
|
||||
assert len(created) == 1
|
||||
server = created[0]
|
||||
with mcp_tool._lock:
|
||||
assert mcp_tool._servers["initial-failure"] is server
|
||||
assert "deterministic initial failure" in (
|
||||
mcp_tool._server_connect_errors["initial-failure"]
|
||||
)
|
||||
assert server._task is not None
|
||||
assert not server._task.done(), "recoverable initial failure was not parked"
|
||||
|
||||
mcp_tool.shutdown_mcp_servers()
|
||||
|
||||
assert pending_at_stop == [], (
|
||||
"shutdown left MCP tasks pending at loop stop: "
|
||||
f"{pending_at_stop!r}"
|
||||
)
|
||||
assert server._task.done()
|
||||
with mcp_tool._lock:
|
||||
assert mcp_tool._mcp_loop is None
|
||||
assert mcp_tool._mcp_thread is None
|
||||
finally:
|
||||
monkeypatch.setattr(mcp_tool, "_stop_mcp_loop", real_stop)
|
||||
_cleanup_mcp_state(mcp_tool, created)
|
||||
|
||||
|
||||
def test_initial_connect_failure_revives_same_registered_server(monkeypatch, tmp_path):
|
||||
"""A cached parked failure must revive through register_mcp_servers()."""
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
||||
|
||||
from tools import mcp_tool
|
||||
from tools.registry import ToolRegistry
|
||||
import tools.registry as registry_module
|
||||
|
||||
_reset_mcp_state(mcp_tool)
|
||||
created = []
|
||||
backend_up = threading.Event()
|
||||
revived = threading.Event()
|
||||
state = {"transport_calls": 0, "tool_calls": 0}
|
||||
mock_registry = ToolRegistry()
|
||||
|
||||
class _Session:
|
||||
async def call_tool(self, name, arguments):
|
||||
state["tool_calls"] += 1
|
||||
return SimpleNamespace(
|
||||
isError=False,
|
||||
content=[SimpleNamespace(text=f"revived:{arguments['value']}")],
|
||||
structuredContent=None,
|
||||
)
|
||||
|
||||
class _RecoveringServerTask(mcp_tool.MCPServerTask):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
created.append(self)
|
||||
|
||||
async def _run_stdio(self, config):
|
||||
assert mcp_tool._connect_server_claim.get() is None
|
||||
state["transport_calls"] += 1
|
||||
if not backend_up.is_set():
|
||||
raise ConnectionError("backend still booting")
|
||||
|
||||
self.session = _Session()
|
||||
self._tools = [SimpleNamespace(
|
||||
name="ping",
|
||||
description="Return a deterministic revival result",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {"value": {"type": "string"}},
|
||||
"required": ["value"],
|
||||
},
|
||||
)]
|
||||
# Match the real transports: discovery runs before _ready is set.
|
||||
self._register_discovered_tools_if_needed()
|
||||
self._ready.set()
|
||||
revived.set()
|
||||
return await self._wait_for_lifecycle_event()
|
||||
|
||||
monkeypatch.setattr(mcp_tool, "MCPServerTask", _RecoveringServerTask)
|
||||
monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True)
|
||||
monkeypatch.setattr(mcp_tool, "_MAX_INITIAL_CONNECT_RETRIES", 0)
|
||||
monkeypatch.setattr(mcp_tool, "_PARKED_RETRY_INTERVAL", 3600)
|
||||
monkeypatch.setattr(registry_module, "registry", mock_registry)
|
||||
|
||||
config = {
|
||||
"recovering": {"command": "unused", "connect_timeout": 5}
|
||||
}
|
||||
|
||||
try:
|
||||
assert mcp_tool.register_mcp_servers(config) == []
|
||||
assert len(created) == 1
|
||||
server = created[0]
|
||||
with mcp_tool._lock:
|
||||
assert mcp_tool._servers["recovering"] is server
|
||||
assert "backend still booting" in (
|
||||
mcp_tool._server_connect_errors["recovering"]
|
||||
)
|
||||
assert not server._task.done()
|
||||
|
||||
backend_up.set()
|
||||
mcp_tool.register_mcp_servers(config)
|
||||
|
||||
assert revived.wait(timeout=5), "cached parked server did not revive"
|
||||
assert len(created) == 1, "revival created a duplicate server task"
|
||||
with mcp_tool._lock:
|
||||
assert mcp_tool._servers["recovering"] is server
|
||||
assert "recovering" not in mcp_tool._server_connect_errors
|
||||
assert state["transport_calls"] == 2
|
||||
assert server.session is not None
|
||||
assert server._error is None
|
||||
|
||||
entry = mock_registry.get_entry("mcp__recovering__ping")
|
||||
assert entry is not None
|
||||
assert entry.check_fn() is True
|
||||
assert json.loads(entry.handler({"value": "ok"})) == {
|
||||
"result": "revived:ok"
|
||||
}
|
||||
assert state["tool_calls"] == 1
|
||||
finally:
|
||||
_cleanup_mcp_state(mcp_tool, created)
|
||||
|
||||
|
||||
def test_terminal_initial_failure_is_not_retained(monkeypatch, tmp_path):
|
||||
"""A non-recoverable startup error must not leave a dead cache entry."""
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
||||
|
||||
from tools import mcp_tool
|
||||
|
||||
_reset_mcp_state(mcp_tool)
|
||||
created = []
|
||||
|
||||
class _AuthFailingServerTask(mcp_tool.MCPServerTask):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
created.append(self)
|
||||
|
||||
async def _run_stdio(self, config):
|
||||
raise PermissionError("terminal authentication failure")
|
||||
|
||||
monkeypatch.setattr(mcp_tool, "MCPServerTask", _AuthFailingServerTask)
|
||||
monkeypatch.setattr(mcp_tool, "_MCP_AVAILABLE", True)
|
||||
monkeypatch.setattr(mcp_tool, "_is_auth_error", lambda exc: True)
|
||||
|
||||
try:
|
||||
assert mcp_tool.register_mcp_servers({
|
||||
"auth-failure": {"command": "unused", "connect_timeout": 5}
|
||||
}) == []
|
||||
assert len(created) == 1
|
||||
assert created[0]._task.done()
|
||||
with mcp_tool._lock:
|
||||
assert "auth-failure" not in mcp_tool._servers
|
||||
assert "terminal authentication failure" in (
|
||||
mcp_tool._server_connect_errors["auth-failure"]
|
||||
)
|
||||
finally:
|
||||
_cleanup_mcp_state(mcp_tool, created)
|
||||
|
||||
|
||||
def test_standalone_failed_connect_is_reaped_without_global_owner(monkeypatch, tmp_path):
|
||||
"""Probe-only _connect_server failures must not publish parked servers."""
|
||||
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
|
||||
|
||||
from tools import mcp_tool
|
||||
|
||||
_reset_mcp_state(mcp_tool)
|
||||
created = []
|
||||
|
||||
class _ProbeServerTask(mcp_tool.MCPServerTask):
|
||||
def __init__(self, name):
|
||||
super().__init__(name)
|
||||
created.append(self)
|
||||
|
||||
async def _run_stdio(self, config):
|
||||
raise ConnectionError("probe target unavailable")
|
||||
|
||||
monkeypatch.setattr(mcp_tool, "MCPServerTask", _ProbeServerTask)
|
||||
monkeypatch.setattr(mcp_tool, "_MAX_INITIAL_CONNECT_RETRIES", 0)
|
||||
monkeypatch.setattr(mcp_tool, "_PARKED_RETRY_INTERVAL", 3600)
|
||||
mcp_tool._ensure_mcp_loop()
|
||||
|
||||
try:
|
||||
with pytest.raises(ConnectionError, match="probe target unavailable"):
|
||||
mcp_tool._run_on_mcp_loop(
|
||||
lambda: mcp_tool._connect_server(
|
||||
"probe-only", {"command": "unused"}
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
assert len(created) == 1
|
||||
assert created[0]._task.done()
|
||||
with mcp_tool._lock:
|
||||
assert "probe-only" not in mcp_tool._servers
|
||||
assert "probe-only" not in mcp_tool._server_connect_errors
|
||||
finally:
|
||||
_cleanup_mcp_state(mcp_tool, created)
|
||||
|
|
@ -514,3 +514,197 @@ class TestMCPInitialConnectionRetry:
|
|||
await task
|
||||
|
||||
asyncio.get_event_loop().run_until_complete(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fix: drain pending tasks before closing the MCP loop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestMCPLoopDrainOnStop:
|
||||
"""_stop_mcp_loop reaps pending tasks while the loop is still open."""
|
||||
|
||||
def test_pending_task_cleanup_runs_before_close(self):
|
||||
"""A task left on the loop must finish its cleanup before close.
|
||||
|
||||
Regression for #60197: a task still suspended when the loop closes is
|
||||
resumed later by the GC, and its ``finally`` then drives ``cancel()``
|
||||
-> ``call_soon()`` against the closed loop, surfacing as an ignored
|
||||
``RuntimeError: Event loop is closed``. Servers absent from
|
||||
``_servers`` (e.g. parked after an initial-connect failure) never get
|
||||
``shutdown()``, so this drain is their only reaper.
|
||||
"""
|
||||
import time
|
||||
import tools.mcp_tool as mcp_mod
|
||||
|
||||
state = {
|
||||
"started": False,
|
||||
"cleanup_ran": False,
|
||||
"cleanup_error": None,
|
||||
"task": None,
|
||||
}
|
||||
|
||||
async def _parked():
|
||||
state["task"] = asyncio.current_task()
|
||||
state["started"] = True
|
||||
try:
|
||||
await asyncio.sleep(3600)
|
||||
finally:
|
||||
# Mirrors _wait_for_reconnect_or_shutdown's finally: cancelling
|
||||
# a helper task needs call_soon(), which a closed loop rejects.
|
||||
try:
|
||||
helper = asyncio.ensure_future(asyncio.sleep(0))
|
||||
helper.cancel()
|
||||
state["cleanup_ran"] = True
|
||||
except BaseException as exc:
|
||||
state["cleanup_error"] = exc
|
||||
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers.clear()
|
||||
mcp_mod._server_connecting.clear()
|
||||
try:
|
||||
mcp_mod._ensure_mcp_loop()
|
||||
with mcp_mod._lock:
|
||||
loop = mcp_mod._mcp_loop
|
||||
assert loop is not None
|
||||
asyncio.run_coroutine_threadsafe(_parked(), loop)
|
||||
|
||||
deadline = time.monotonic() + 5
|
||||
while not state["started"] and time.monotonic() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert state["started"], "task never started on the MCP loop"
|
||||
|
||||
stop_saw_cleanup = []
|
||||
call_soon = loop.call_soon
|
||||
call_soon_threadsafe = loop.call_soon_threadsafe
|
||||
|
||||
def record_stop_order(schedule, callback, *args, **kwargs):
|
||||
if (
|
||||
getattr(callback, "__self__", None) is loop
|
||||
and getattr(callback, "__name__", None) == "stop"
|
||||
):
|
||||
stop_saw_cleanup.append(state["cleanup_ran"])
|
||||
return schedule(callback, *args, **kwargs)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
loop,
|
||||
"call_soon",
|
||||
side_effect=lambda callback, *args, **kwargs: record_stop_order(
|
||||
call_soon, callback, *args, **kwargs
|
||||
),
|
||||
),
|
||||
patch.object(
|
||||
loop,
|
||||
"call_soon_threadsafe",
|
||||
side_effect=lambda callback, *args, **kwargs: record_stop_order(
|
||||
call_soon_threadsafe, callback, *args, **kwargs
|
||||
),
|
||||
),
|
||||
):
|
||||
mcp_mod._stop_mcp_loop()
|
||||
|
||||
assert state["task"] is not None
|
||||
assert state["task"].done(), "task left pending when the loop closed"
|
||||
assert state["cleanup_error"] is None, (
|
||||
f"cleanup ran against a closed loop: {state['cleanup_error']!r}"
|
||||
)
|
||||
assert state["cleanup_ran"], "task cleanup never ran"
|
||||
assert stop_saw_cleanup == [True], "loop.stop ran before task cleanup"
|
||||
finally:
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers.clear()
|
||||
mcp_mod._server_connecting.clear()
|
||||
mcp_mod._stop_mcp_loop()
|
||||
|
||||
def test_drain_is_bounded_when_task_ignores_cancellation(self, caplog):
|
||||
"""A cancellation-resistant task must not hang final MCP shutdown."""
|
||||
import tools.mcp_tool as mcp_mod
|
||||
|
||||
async def _run():
|
||||
release = asyncio.Event()
|
||||
|
||||
async def cancellation_resistant():
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
except asyncio.CancelledError:
|
||||
await release.wait()
|
||||
|
||||
task = asyncio.create_task(cancellation_resistant())
|
||||
await asyncio.sleep(0)
|
||||
try:
|
||||
with caplog.at_level("WARNING", logger=mcp_mod.logger.name):
|
||||
async with asyncio.timeout(0.5):
|
||||
await mcp_mod._drain_mcp_loop_tasks(timeout=0.01)
|
||||
assert not task.done(), "drain waited indefinitely for resistant task"
|
||||
finally:
|
||||
release.set()
|
||||
if not task.done() and task.cancelling() == 0:
|
||||
task.cancel()
|
||||
await task
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
assert any(
|
||||
"still pending after" in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
def test_outer_timeout_still_allows_loop_owned_drain_before_stop(
|
||||
self, caplog, monkeypatch
|
||||
):
|
||||
"""A blocked loop must drain after it resumes, not stop ahead of the drain."""
|
||||
import threading
|
||||
import tools.mcp_tool as mcp_mod
|
||||
|
||||
parked_started = threading.Event()
|
||||
cleanup_ran = threading.Event()
|
||||
blocker_started = threading.Event()
|
||||
release_blocker = threading.Event()
|
||||
|
||||
async def parked_task():
|
||||
parked_started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cleanup_ran.set()
|
||||
|
||||
def block_loop():
|
||||
blocker_started.set()
|
||||
release_blocker.wait(timeout=5)
|
||||
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers.clear()
|
||||
mcp_mod._server_connecting.clear()
|
||||
mcp_mod._ensure_mcp_loop()
|
||||
with mcp_mod._lock:
|
||||
loop = mcp_mod._mcp_loop
|
||||
assert loop is not None
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(parked_task(), loop)
|
||||
assert parked_started.wait(timeout=2)
|
||||
loop.call_soon_threadsafe(block_loop)
|
||||
assert blocker_started.wait(timeout=2)
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "_MCP_LOOP_DRAIN_TIMEOUT", 0.01)
|
||||
release_timer = threading.Timer(1.2, release_blocker.set)
|
||||
release_timer.start()
|
||||
try:
|
||||
with caplog.at_level("WARNING", logger=mcp_mod.logger.name):
|
||||
mcp_mod._stop_mcp_loop()
|
||||
|
||||
assert cleanup_ran.is_set(), "drain was overtaken by loop.stop"
|
||||
assert future.done(), "parked task remained pending after loop resumed"
|
||||
assert loop.is_closed()
|
||||
finally:
|
||||
release_timer.cancel()
|
||||
release_blocker.set()
|
||||
future.cancel()
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers.clear()
|
||||
mcp_mod._server_connecting.clear()
|
||||
mcp_mod._stop_mcp_loop()
|
||||
|
||||
assert any(
|
||||
"Timed out waiting for MCP loop drain" in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
|
|
|||
|
|
@ -888,6 +888,86 @@ class TestGracefulFallback:
|
|||
|
||||
class TestShutdown:
|
||||
|
||||
def test_shutdown_drains_parked_server_after_bounded_wait_expires(self):
|
||||
"""The public shutdown path drains a parked server if graceful shutdown stalls.
|
||||
|
||||
This exercises the production ownership path: a real ``MCPServerTask``
|
||||
is registered, its parked waiter owns child tasks on the shared loop,
|
||||
and the bounded wait for the scheduled shutdown expires. The loop owner
|
||||
must still cancel and drain that waiter before closing the loop.
|
||||
"""
|
||||
import tools.mcp_tool as mcp_mod
|
||||
from tools.mcp_tool import MCPServerTask, shutdown_mcp_servers
|
||||
|
||||
shutdown_started = threading.Event()
|
||||
parked_task_done = threading.Event()
|
||||
scheduled_shutdown = {}
|
||||
schedule_count = 0
|
||||
|
||||
class StalledShutdownServer(MCPServerTask):
|
||||
async def shutdown(self):
|
||||
shutdown_started.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
server = StalledShutdownServer("parked")
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers.clear()
|
||||
mcp_mod._server_connecting.clear()
|
||||
mcp_mod._ensure_mcp_loop()
|
||||
with mcp_mod._lock:
|
||||
loop = mcp_mod._mcp_loop
|
||||
assert loop is not None
|
||||
|
||||
async def install_parked_waiter():
|
||||
task = asyncio.create_task(server._wait_for_reconnect_or_shutdown())
|
||||
server._task = task
|
||||
task.add_done_callback(lambda _task: parked_task_done.set())
|
||||
await asyncio.sleep(0)
|
||||
return task
|
||||
|
||||
parked_task = asyncio.run_coroutine_threadsafe(
|
||||
install_parked_waiter(), loop
|
||||
).result(timeout=2)
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers[server.name] = server
|
||||
|
||||
def schedule_then_report_timeout(coro, target_loop, **_kwargs):
|
||||
nonlocal schedule_count
|
||||
schedule_count += 1
|
||||
future = asyncio.run_coroutine_threadsafe(coro, target_loop)
|
||||
if schedule_count > 1:
|
||||
return future
|
||||
|
||||
scheduled_shutdown["future"] = future
|
||||
|
||||
class TimedOutFuture:
|
||||
def result(self, timeout):
|
||||
assert timeout > 0
|
||||
assert shutdown_started.wait(timeout=2)
|
||||
raise TimeoutError("simulated bounded MCP shutdown timeout")
|
||||
|
||||
return TimedOutFuture()
|
||||
|
||||
try:
|
||||
with patch(
|
||||
"agent.async_utils.safe_schedule_threadsafe",
|
||||
side_effect=schedule_then_report_timeout,
|
||||
):
|
||||
shutdown_mcp_servers()
|
||||
|
||||
assert loop.is_closed()
|
||||
assert parked_task_done.is_set(), (
|
||||
"parked MCPServerTask was not drained before its loop closed"
|
||||
)
|
||||
assert parked_task.done()
|
||||
assert scheduled_shutdown["future"].done()
|
||||
assert schedule_count == 2
|
||||
finally:
|
||||
with mcp_mod._lock:
|
||||
mcp_mod._servers.clear()
|
||||
mcp_mod._server_connecting.clear()
|
||||
mcp_mod._stop_mcp_loop()
|
||||
|
||||
def test_shutdown_deregisters_registered_tools(self):
|
||||
"""shutdown_mcp_servers removes MCP tools and their raw alias."""
|
||||
import tools.mcp_tool as mcp_mod
|
||||
|
|
|
|||
|
|
@ -363,6 +363,11 @@ def _jittered(seconds: float) -> float:
|
|||
_DEFAULT_KEEPALIVE_INTERVAL = 180 # seconds between liveness pings
|
||||
_MIN_KEEPALIVE_INTERVAL = 5 # clamp floor for configured intervals
|
||||
|
||||
# Final shutdown gives pending MCP-loop tasks one bounded cancellation cycle
|
||||
# before closing their owning loop. Cooperative parked/reconnect waiters finish
|
||||
# immediately; cancellation-resistant tasks must not hang process exit.
|
||||
_MCP_LOOP_DRAIN_TIMEOUT = 3.0
|
||||
|
||||
# Environment variables that are safe to pass to stdio subprocesses
|
||||
_SAFE_ENV_KEYS = frozenset({
|
||||
"PATH", "HOME", "USER", "LANG", "LC_ALL", "TERM", "SHELL", "TMPDIR",
|
||||
|
|
@ -3016,7 +3021,7 @@ class MCPServerTask:
|
|||
self._register_discovered_tools_if_needed()
|
||||
|
||||
def _register_discovered_tools_if_needed(self) -> None:
|
||||
"""Re-register tools after a post-ready reconnect if needed.
|
||||
"""Re-register tools after an owned server reconnects if needed.
|
||||
|
||||
Initial registration is performed by ``_discover_and_register_server``
|
||||
after ``start()`` completes. During a later reconnect, outage handling
|
||||
|
|
@ -3024,6 +3029,9 @@ class MCPServerTask:
|
|||
A managed server can still be identified by its entry in ``_servers``;
|
||||
publish its freshly discovered tools before transport readiness is
|
||||
restored so a successful revival cannot come back with zero tools.
|
||||
A server retained after a recoverable initial failure is likewise
|
||||
registry-owned before its first successful session, so ownership also
|
||||
authorizes its first publication.
|
||||
"""
|
||||
if self._registered_tool_names:
|
||||
return
|
||||
|
|
@ -3034,6 +3042,12 @@ class MCPServerTask:
|
|||
self._registered_tool_names = _register_server_tools(
|
||||
self.name, self, self._config
|
||||
)
|
||||
# A retained initial-failure server that just published tools has
|
||||
# recovered: drop its stale connect error so status surfaces stop
|
||||
# reporting it as failed.
|
||||
with _lock:
|
||||
if _servers.get(self.name) is self:
|
||||
_server_connect_errors.pop(self.name, None)
|
||||
|
||||
async def run(self, config: dict):
|
||||
"""Long-lived coroutine: connect, discover tools, wait, disconnect.
|
||||
|
|
@ -3518,6 +3532,12 @@ class MCPServerTask:
|
|||
_servers: Dict[str, MCPServerTask] = {}
|
||||
_server_connecting: set[str] = set()
|
||||
_server_connect_errors: Dict[str, str] = {}
|
||||
# Discovery installs a task-local claim before calling ``_connect_server`` so
|
||||
# it can retain a recoverable parked task without making standalone probe calls
|
||||
# publish failed servers into module-global ownership.
|
||||
_connect_server_claim: contextvars.ContextVar[
|
||||
Optional[Callable[[MCPServerTask], None]]
|
||||
] = contextvars.ContextVar("mcp_connect_server_claim", default=None)
|
||||
|
||||
# Connection-retry cooldown (per-server isolation against restart storms).
|
||||
#
|
||||
|
|
@ -3627,9 +3647,8 @@ def _signal_reconnect(server: Any) -> bool:
|
|||
The tool handlers run on caller threads, while the server task and its
|
||||
``_reconnect_event`` live on the background MCP loop. Setting an
|
||||
asyncio.Event from another thread must go through
|
||||
``loop.call_soon_threadsafe``; only fall back to a direct ``.set()``
|
||||
when the loop isn't running (e.g. unit tests that drive the handler
|
||||
synchronously).
|
||||
``loop.call_soon_threadsafe``; non-async adapters and tests without a
|
||||
running loop can use a direct ``.set()``.
|
||||
|
||||
Returns True if a reconnect signal was delivered, False if the server
|
||||
has no reconnect machinery (nothing to revive).
|
||||
|
|
@ -3638,7 +3657,11 @@ def _signal_reconnect(server: Any) -> bool:
|
|||
if event is None:
|
||||
return False
|
||||
loop = _mcp_loop
|
||||
if loop is not None and loop.is_running():
|
||||
if (
|
||||
isinstance(event, asyncio.Event)
|
||||
and loop is not None
|
||||
and loop.is_running()
|
||||
):
|
||||
loop.call_soon_threadsafe(event.set)
|
||||
else:
|
||||
event.set()
|
||||
|
|
@ -4611,7 +4634,37 @@ async def _connect_server(name: str, config: dict) -> MCPServerTask:
|
|||
Exception: on connection or initialization failure.
|
||||
"""
|
||||
server = MCPServerTask(name)
|
||||
await server.start(config)
|
||||
claim = _connect_server_claim.get()
|
||||
claim_token = None
|
||||
if claim is not None:
|
||||
claim(server)
|
||||
# ``start()`` creates the long-lived run task by copying this context.
|
||||
# The ownership callback is only for this connection attempt; do not
|
||||
# retain its discovery closure for the server's lifetime.
|
||||
claim_token = _connect_server_claim.set(None)
|
||||
try:
|
||||
await server.start(config)
|
||||
except asyncio.CancelledError:
|
||||
# start() already cancels/reaps server._task on external cancellation
|
||||
# (see the comment there) -- awaiting a redundant shutdown() inside a
|
||||
# cancelled context would only risk swallowing the cancellation.
|
||||
raise
|
||||
except BaseException:
|
||||
# Discovery owns claimed tasks and decides whether a failed start is a
|
||||
# live recoverable park or a terminal failure. Standalone probes have
|
||||
# no revival owner, so they must reap their failed task locally.
|
||||
if claim is None:
|
||||
try:
|
||||
await server.shutdown()
|
||||
except Exception as shutdown_exc: # noqa: BLE001 -- best-effort reap, don't mask the real error
|
||||
logger.debug(
|
||||
"MCP server '%s' shutdown during orphan-reap failed: %s",
|
||||
name, shutdown_exc,
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
if claim_token is not None:
|
||||
_connect_server_claim.reset(claim_token)
|
||||
return server
|
||||
|
||||
|
||||
|
|
@ -5783,10 +5836,45 @@ async def _discover_and_register_server(name: str, config: dict) -> List[str]:
|
|||
Returns list of registered tool names.
|
||||
"""
|
||||
connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT)
|
||||
server = await asyncio.wait_for(
|
||||
_connect_server(name, config),
|
||||
timeout=connect_timeout,
|
||||
)
|
||||
# List-based claim (not a ``nonlocal`` rebind): the claim callback runs
|
||||
# inside ``_connect_server`` while this frame is suspended, and appending
|
||||
# keeps type narrowing intact for the module's other ``server`` locals.
|
||||
claimed: List[MCPServerTask] = []
|
||||
|
||||
def _claim_server(created: MCPServerTask) -> None:
|
||||
claimed.append(created)
|
||||
|
||||
claim_token = _connect_server_claim.set(_claim_server)
|
||||
try:
|
||||
server = await asyncio.wait_for(
|
||||
_connect_server(name, config),
|
||||
timeout=connect_timeout,
|
||||
)
|
||||
except BaseException:
|
||||
server = claimed[0] if claimed else None
|
||||
task = server._task if server is not None else None
|
||||
task_cancelling = (
|
||||
task.cancelling()
|
||||
if task is not None and hasattr(task, "cancelling")
|
||||
else 0
|
||||
)
|
||||
if (
|
||||
server is not None
|
||||
and server._error is not None
|
||||
and task is not None
|
||||
and not task.done()
|
||||
and not task_cancelling
|
||||
):
|
||||
# Recoverable park: the run task deliberately stays alive to
|
||||
# self-probe, so adopt it into the registry for shutdown/revival.
|
||||
with _lock:
|
||||
_servers[name] = server
|
||||
elif server is not None:
|
||||
await server.shutdown()
|
||||
raise
|
||||
finally:
|
||||
_connect_server_claim.reset(claim_token)
|
||||
|
||||
with _lock:
|
||||
_server_connecting.discard(name)
|
||||
_server_connect_errors.pop(name, None)
|
||||
|
|
@ -5927,7 +6015,11 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
|
|||
|
||||
# Log a summary so ACP callers get visibility into what was registered.
|
||||
with _lock:
|
||||
connected = [n for n in new_servers if n in _servers]
|
||||
connected = [
|
||||
n
|
||||
for n in new_servers
|
||||
if n in _servers and n not in _server_connect_errors
|
||||
]
|
||||
new_tool_count = sum(
|
||||
len(getattr(_servers[n], "_registered_tool_names", []))
|
||||
for n in connected
|
||||
|
|
@ -6000,7 +6092,11 @@ def discover_mcp_tools() -> List[str]:
|
|||
return tool_names
|
||||
|
||||
with _lock:
|
||||
connected_server_names = [name for name in new_server_names if name in _servers]
|
||||
connected_server_names = [
|
||||
name
|
||||
for name in new_server_names
|
||||
if name in _servers and name not in _server_connect_errors
|
||||
]
|
||||
new_tool_count = sum(
|
||||
len(getattr(_servers[name], "_registered_tool_names", []))
|
||||
for name in connected_server_names
|
||||
|
|
@ -6618,6 +6714,58 @@ def _stop_mcp_loop_if_idle() -> bool:
|
|||
return _stop_mcp_loop(only_if_idle=True)
|
||||
|
||||
|
||||
async def _drain_mcp_loop_tasks(
|
||||
*,
|
||||
timeout: float = _MCP_LOOP_DRAIN_TIMEOUT,
|
||||
) -> None:
|
||||
"""Cancel every task still pending on the MCP loop and reap it.
|
||||
|
||||
Cancelling is not enough on its own: ``Task.cancel()`` only schedules the
|
||||
throw, so tasks need a cancellation cycle before the loop goes away. Wait
|
||||
for them here — on their owning loop — but keep the final drain bounded so
|
||||
a task that suppresses cancellation cannot hang process exit indefinitely.
|
||||
"""
|
||||
current = asyncio.current_task()
|
||||
pending = [t for t in asyncio.all_tasks() if t is not current and not t.done()]
|
||||
if not pending:
|
||||
return
|
||||
logger.debug("Draining %d pending task(s) from the MCP loop", len(pending))
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
|
||||
done, still_pending = await asyncio.wait(pending, timeout=timeout)
|
||||
for task in done:
|
||||
if task.cancelled():
|
||||
continue
|
||||
try:
|
||||
task.exception()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.debug("Pending MCP loop task ended during shutdown: %s", exc)
|
||||
|
||||
if still_pending:
|
||||
logger.warning(
|
||||
"%d MCP loop task(s) still pending after %.1fs drain",
|
||||
len(still_pending), timeout,
|
||||
)
|
||||
|
||||
|
||||
async def _drain_and_stop_mcp_loop() -> None:
|
||||
"""Drain pending tasks, then stop the loop from its owning thread.
|
||||
|
||||
Keeping both operations in one loop-owned sequence matters when the caller
|
||||
times out waiting for a blocked loop. Queuing ``loop.stop`` separately from
|
||||
the caller can overtake the scheduled drain before it receives a loop cycle,
|
||||
leaving the drain coroutine itself pending when the loop is closed.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
await _drain_mcp_loop_tasks(timeout=_MCP_LOOP_DRAIN_TIMEOUT)
|
||||
finally:
|
||||
loop.call_soon(loop.stop)
|
||||
|
||||
|
||||
def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool:
|
||||
"""Stop the background event loop and join its thread."""
|
||||
global _mcp_loop, _mcp_thread
|
||||
|
|
@ -6630,13 +6778,50 @@ def _stop_mcp_loop(*, only_if_idle: bool = False) -> bool:
|
|||
_mcp_loop = None
|
||||
_mcp_thread = None
|
||||
if loop is not None:
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
# Drain before stopping: closing the loop with tasks still suspended
|
||||
# leaves their coroutines for the GC, whose finalizer then resumes them
|
||||
# to run cleanup against a loop that is already closed -> "Event loop
|
||||
# is closed" (#60197). ``shutdown_mcp_servers`` only reaps servers held
|
||||
# in ``_servers``, so anything else left on this loop ends up here.
|
||||
stop_owned_by_loop = False
|
||||
if loop.is_running():
|
||||
from agent.async_utils import safe_schedule_threadsafe
|
||||
|
||||
future = safe_schedule_threadsafe(
|
||||
_drain_and_stop_mcp_loop(), loop,
|
||||
logger=logger,
|
||||
log_message="MCP loop drain: failed to schedule",
|
||||
log_level=logging.WARNING,
|
||||
)
|
||||
if future is not None:
|
||||
stop_owned_by_loop = True
|
||||
try:
|
||||
future.result(timeout=_MCP_LOOP_DRAIN_TIMEOUT + 1)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"Timed out waiting for MCP loop drain after %.1fs",
|
||||
_MCP_LOOP_DRAIN_TIMEOUT + 1,
|
||||
)
|
||||
except BaseException as exc:
|
||||
logger.warning("Error draining MCP loop tasks: %s", exc)
|
||||
elif not loop.is_closed():
|
||||
try:
|
||||
loop.run_until_complete(
|
||||
_drain_mcp_loop_tasks(timeout=_MCP_LOOP_DRAIN_TIMEOUT)
|
||||
)
|
||||
except BaseException as exc:
|
||||
logger.warning("Error draining stopped MCP loop tasks: %s", exc)
|
||||
|
||||
if not stop_owned_by_loop and loop.is_running():
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
if thread is not None:
|
||||
thread.join(timeout=5)
|
||||
if thread.is_alive():
|
||||
logger.warning("MCP event loop thread did not stop within 5.0s")
|
||||
try:
|
||||
loop.close()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.warning("Unable to close MCP event loop cleanly: %s", exc)
|
||||
# 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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue