mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-30 19:09:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

fix(mcp): retain parked startup tasks for clean shutdown

Browse source
This commit is contained in:
Stepan Zadolia 2026-07-10 13:21:39 +03:00 • committed by kshitij
parent fd5397d69a
commit c00a1d58d5
2 changed files with 362 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

273
tests/tools/test_mcp_initial_connect_shutdown.py Normal file
View 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)

101
tools/mcp_tool.py
View file

@ -3021,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
@ -3029,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
@ -3039,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.
@ -3523,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).
#
@ -3632,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).
@ -3643,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()
@ -4616,7 +4634,26 @@ 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 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:
await server.shutdown()
raise
finally:
if claim_token is not None:
_connect_server_claim.reset(claim_token)
return server
@ -5788,10 +5825,42 @@ 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,
)
server: Optional[MCPServerTask] = None
def _claim_server(created: MCPServerTask) -> None:
nonlocal server
server = created
claim_token = _connect_server_claim.set(_claim_server)
try:
server = await asyncio.wait_for(
_connect_server(name, config),
timeout=connect_timeout,
)
except BaseException:
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
)
recoverable_park = (
server is not None
and server._error is not None
and task is not None
and not task.done()
and not task_cancelling
)
if recoverable_park:
with _lock:
_servers[name] = server
elif server is not None:
await server.shutdown()
raise
finally:
_connect_server_claim.reset(claim_token)
assert server is not None
with _lock:
_server_connecting.discard(name)
_server_connect_errors.pop(name, None)
@ -5932,7 +6001,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
@ -6005,7 +6078,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