fix(mcp): re-register tools during parked revival

This commit is contained in:
yungchentang 2026-07-19 02:13:18 +02:00 committed by Teknium
parent da8b89c835
commit 106d1822e3
2 changed files with 39 additions and 6 deletions

View file

@ -8,10 +8,39 @@ revival probe on its own.
"""
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
def test_revival_discovery_registers_tools_while_ready_is_cleared(monkeypatch):
"""A managed server revival must publish tools before readiness is reset."""
from tools import mcp_tool
from tools.mcp_tool import MCPServerTask
server = MCPServerTask("srv")
server._config = {"url": "https://example.test/mcp"}
server.session = SimpleNamespace(
list_tools=AsyncMock(
return_value=SimpleNamespace(
tools=[SimpleNamespace(name="send_message")],
)
)
)
server._ready.clear()
server._registered_tool_names = []
monkeypatch.setitem(mcp_tool._servers, server.name, server)
register = MagicMock(return_value=["srv__send_message"])
monkeypatch.setattr(mcp_tool, "_register_server_tools", register)
asyncio.run(server._discover_tools())
register.assert_called_once_with(server.name, server, server._config)
assert server._registered_tool_names == ["srv__send_message"]
@pytest.mark.no_isolate
def test_parked_server_self_probes_and_revives(monkeypatch, tmp_path):
"""A parked server must revive on its own once the backend recovers,

View file

@ -2986,14 +2986,18 @@ class MCPServerTask:
"""Re-register tools after a post-ready reconnect if needed.
Initial registration is performed by ``_discover_and_register_server``
after ``start()`` completes. During a later reconnect, however,
``_ready`` remains set; if outage handling previously deregistered
stale tools (parking calls ``_deregister_tools``), a successful
revival must publish the freshly discovered tools again otherwise
the transport comes back alive with zero registered tools.
after ``start()`` completes. During a later reconnect, outage handling
may clear ``_ready`` before discovery and may deregister stale tools.
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.
"""
if not self._ready.is_set() or self._registered_tool_names:
if self._registered_tool_names:
return
if not self._ready.is_set():
with _lock:
if _servers.get(self.name) is not self:
return
self._registered_tool_names = _register_server_tools(
self.name, self, self._config
)