feat(mcp): lazy server startup from schema cache (design from #56832)
Wires the fingerprint-keyed schema cache (previous commit, @Vansh5632's design from #56832) into the startup path, re-derived onto main's current connect machinery: - register_mcp_servers: servers with mcp_servers.<name>.lazy=true whose config fingerprint matches a valid cache entry register tools from cache WITHOUT spawning; miss/stale falls back to eager connect. - First tool use routes through _ensure_lazy_server_connected, which composes with the connect cooldown (#50394) and _server_connecting dedup rather than duplicating the connect path. - resource/prompt utility handlers (list_resources/get_prompt) also connect-on-first-use — closes the gap flagged in the original sweeper review. - Write-through: a live connect refreshes the cache entry. Config gate is per-server, default OFF, matching the idle_timeout_seconds key pattern. 24 lazy/cache tests + 440 mcp-wide green; mutation-checked (cache-read disabled -> registration test fails; connect bypassed -> 3 first-use tests fail).
This commit is contained in:
parent
135a29452a
commit
1d5ecad568
|
|
@ -0,0 +1,281 @@
|
|||
"""Behavior-contract tests for lazy MCP server startup (#56832).
|
||||
|
||||
A server configured with ``lazy: true`` whose config fingerprint matches an
|
||||
on-disk schema-cache entry registers its tools WITHOUT spawning/connecting;
|
||||
the first real call (raw tool OR resource/prompt utility) routes through the
|
||||
existing connect path.
|
||||
"""
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import tools.mcp_tool as mcp
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_mcp_state():
|
||||
old_servers = dict(mcp._servers)
|
||||
old_lazy = dict(mcp._lazy_server_configs)
|
||||
old_fps = dict(mcp._lazy_server_fingerprints)
|
||||
old_names = dict(mcp._lazy_server_tool_names)
|
||||
old_connecting = set(mcp._server_connecting)
|
||||
yield
|
||||
mcp._servers.clear()
|
||||
mcp._servers.update(old_servers)
|
||||
mcp._lazy_server_configs.clear()
|
||||
mcp._lazy_server_configs.update(old_lazy)
|
||||
mcp._lazy_server_fingerprints.clear()
|
||||
mcp._lazy_server_fingerprints.update(old_fps)
|
||||
mcp._lazy_server_tool_names.clear()
|
||||
mcp._lazy_server_tool_names.update(old_names)
|
||||
mcp._server_connecting.clear()
|
||||
mcp._server_connecting.update(old_connecting)
|
||||
|
||||
|
||||
def _fake_cache_entry():
|
||||
return {
|
||||
"fingerprint": "abc",
|
||||
"tools": [
|
||||
{
|
||||
"name": "browser_navigate",
|
||||
"description": "Navigate",
|
||||
"inputSchema": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
"utility_tools": [],
|
||||
}
|
||||
|
||||
|
||||
def _lazy_config():
|
||||
return {
|
||||
"playwright": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "@playwright/mcp"],
|
||||
"lazy": True,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class TestLazyMcpRegistration:
|
||||
def test_registers_from_cache_without_connect(self):
|
||||
config = _lazy_config()
|
||||
with patch("tools.mcp_tool._MCP_AVAILABLE", True), \
|
||||
patch("tools.mcp_schema_cache.config_fingerprint", return_value="abc"), \
|
||||
patch("tools.mcp_schema_cache.get_cached_entry", return_value=_fake_cache_entry()), \
|
||||
patch(
|
||||
"tools.mcp_tool._register_from_cache_sync",
|
||||
return_value=["mcp_playwright_browser_navigate"],
|
||||
) as mock_register, \
|
||||
patch("tools.mcp_tool._discover_and_register_server", new_callable=AsyncMock) as mock_discover, \
|
||||
patch("tools.mcp_tool._ensure_mcp_loop") as mock_loop, \
|
||||
patch("tools.mcp_tool._run_on_mcp_loop") as mock_run:
|
||||
|
||||
mcp.register_mcp_servers(config)
|
||||
|
||||
mock_register.assert_called_once()
|
||||
mock_discover.assert_not_called()
|
||||
mock_run.assert_not_called()
|
||||
mock_loop.assert_not_called()
|
||||
|
||||
def test_cache_miss_falls_back_to_eager_connect(self):
|
||||
config = _lazy_config()
|
||||
with patch("tools.mcp_tool._MCP_AVAILABLE", True), \
|
||||
patch("tools.mcp_schema_cache.config_fingerprint", return_value="abc"), \
|
||||
patch("tools.mcp_schema_cache.get_cached_entry", return_value=None), \
|
||||
patch("tools.mcp_tool._ensure_mcp_loop"), \
|
||||
patch("tools.mcp_tool._run_on_mcp_loop") as mock_run:
|
||||
|
||||
mcp.register_mcp_servers(config)
|
||||
|
||||
mock_run.assert_called_once()
|
||||
|
||||
def test_non_lazy_server_never_touches_cache(self):
|
||||
config = {"playwright": {"command": "npx", "args": []}}
|
||||
with patch("tools.mcp_tool._MCP_AVAILABLE", True), \
|
||||
patch("tools.mcp_schema_cache.get_cached_entry") as mock_get, \
|
||||
patch("tools.mcp_tool._ensure_mcp_loop"), \
|
||||
patch("tools.mcp_tool._run_on_mcp_loop") as mock_run:
|
||||
|
||||
mcp.register_mcp_servers(config)
|
||||
|
||||
mock_get.assert_not_called()
|
||||
mock_run.assert_called_once()
|
||||
|
||||
def test_lazy_server_not_reregistered_on_second_pass(self):
|
||||
config = _lazy_config()
|
||||
mcp._lazy_server_configs["playwright"] = dict(config["playwright"])
|
||||
mcp._lazy_server_tool_names["playwright"] = ["mcp_playwright_browser_navigate"]
|
||||
with patch("tools.mcp_tool._MCP_AVAILABLE", True), \
|
||||
patch("tools.mcp_tool._register_from_cache_sync") as mock_register, \
|
||||
patch("tools.mcp_tool._run_on_mcp_loop") as mock_run:
|
||||
|
||||
names = mcp.register_mcp_servers(config)
|
||||
|
||||
mock_register.assert_not_called()
|
||||
mock_run.assert_not_called()
|
||||
assert "mcp_playwright_browser_navigate" in names
|
||||
|
||||
|
||||
class TestLazyFirstUseConnect:
|
||||
def _connected_server(self):
|
||||
mock_session = MagicMock()
|
||||
mock_session.call_tool = AsyncMock(
|
||||
return_value=SimpleNamespace(isError=False, content=[], structuredContent=None)
|
||||
)
|
||||
connected = SimpleNamespace(
|
||||
session=mock_session,
|
||||
_rpc_lock=MagicMock(),
|
||||
_pending_call_context=None,
|
||||
)
|
||||
connected._rpc_lock.__aenter__ = AsyncMock(return_value=None)
|
||||
connected._rpc_lock.__aexit__ = AsyncMock(return_value=None)
|
||||
return connected
|
||||
|
||||
@staticmethod
|
||||
def _run_on_loop(coro_or_factory, timeout=120):
|
||||
import asyncio
|
||||
|
||||
coro = coro_or_factory() if callable(coro_or_factory) else coro_or_factory
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
return loop.run_until_complete(coro)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
def test_tool_handler_lazy_connects_on_first_call(self):
|
||||
config = {"command": "npx", "args": [], "lazy": True, "timeout": 5}
|
||||
mcp._lazy_server_configs["playwright"] = dict(config)
|
||||
mcp._lazy_server_fingerprints["playwright"] = "abc"
|
||||
|
||||
connected = self._connected_server()
|
||||
|
||||
def _connect(name):
|
||||
mcp._servers["playwright"] = connected
|
||||
return True
|
||||
|
||||
with patch.object(mcp, "_ensure_lazy_server_connected", side_effect=_connect) as mock_connect, \
|
||||
patch.object(mcp, "_run_on_mcp_loop", side_effect=self._run_on_loop):
|
||||
handler = mcp._make_tool_handler("playwright", "browser_navigate", 5)
|
||||
out = handler({}, task_id="t1")
|
||||
|
||||
mock_connect.assert_called_once_with("playwright")
|
||||
payload = json.loads(out)
|
||||
assert "error" not in payload
|
||||
assert payload.get("result") == ""
|
||||
|
||||
def test_list_resources_handler_lazy_connects_on_first_call(self):
|
||||
# Regression for the resource/prompt gap: utility handlers must also
|
||||
# route through the first-use connect path, or the first
|
||||
# list_resources/get_prompt on a lazy server fails.
|
||||
config = {"command": "npx", "args": [], "lazy": True, "timeout": 5}
|
||||
mcp._lazy_server_configs["playwright"] = dict(config)
|
||||
|
||||
connected = self._connected_server()
|
||||
connected.session.list_resources = AsyncMock()
|
||||
|
||||
def _connect(name):
|
||||
mcp._servers["playwright"] = connected
|
||||
return True
|
||||
|
||||
async def _fake_paginate(list_method, items_attr, server_name):
|
||||
return [SimpleNamespace(uri="file:///a", name="a", description="", mimeType="")]
|
||||
|
||||
with patch.object(mcp, "_ensure_lazy_server_connected", side_effect=_connect) as mock_connect, \
|
||||
patch.object(mcp, "_paginate_full_list", side_effect=_fake_paginate), \
|
||||
patch.object(mcp, "_run_on_mcp_loop", side_effect=self._run_on_loop):
|
||||
handler = mcp._make_list_resources_handler("playwright", 5)
|
||||
out = handler({})
|
||||
|
||||
mock_connect.assert_called_once_with("playwright")
|
||||
payload = json.loads(out)
|
||||
assert "error" not in payload
|
||||
assert payload["resources"][0]["uri"] == "file:///a"
|
||||
|
||||
def test_get_prompt_handler_lazy_connects_on_first_call(self):
|
||||
config = {"command": "npx", "args": [], "lazy": True, "timeout": 5}
|
||||
mcp._lazy_server_configs["playwright"] = dict(config)
|
||||
|
||||
connected = self._connected_server()
|
||||
connected.session.get_prompt = AsyncMock(
|
||||
return_value=SimpleNamespace(messages=[])
|
||||
)
|
||||
|
||||
def _connect(name):
|
||||
mcp._servers["playwright"] = connected
|
||||
return True
|
||||
|
||||
with patch.object(mcp, "_ensure_lazy_server_connected", side_effect=_connect) as mock_connect, \
|
||||
patch.object(mcp, "_run_on_mcp_loop", side_effect=self._run_on_loop):
|
||||
handler = mcp._make_get_prompt_handler("playwright", 5)
|
||||
out = handler({"name": "greeting"})
|
||||
|
||||
mock_connect.assert_called_once_with("playwright")
|
||||
payload = json.loads(out)
|
||||
assert "error" not in payload
|
||||
|
||||
def test_check_fn_passes_for_lazy_registered_server(self):
|
||||
mcp._lazy_server_configs["playwright"] = {"lazy": True}
|
||||
mcp._lazy_server_fingerprints["playwright"] = "abc"
|
||||
assert mcp._make_check_fn("playwright")() is True
|
||||
|
||||
def test_check_fn_fails_for_unknown_server(self):
|
||||
assert mcp._make_check_fn("nope")() is False
|
||||
|
||||
def test_lazy_connect_respects_connect_cooldown(self):
|
||||
mcp._lazy_server_configs["playwright"] = {"command": "npx", "lazy": True}
|
||||
with patch.object(mcp, "_connect_cooldown_active", return_value=True), \
|
||||
patch.object(mcp, "_run_on_mcp_loop") as mock_run:
|
||||
assert mcp._ensure_lazy_server_connected("playwright") is False
|
||||
mock_run.assert_not_called()
|
||||
|
||||
def test_lazy_connect_success_clears_lazy_state(self):
|
||||
mcp._lazy_server_configs["playwright"] = {"command": "npx", "lazy": True}
|
||||
mcp._lazy_server_fingerprints["playwright"] = "abc"
|
||||
mcp._lazy_server_tool_names["playwright"] = ["mcp_playwright_browser_navigate"]
|
||||
|
||||
connected = SimpleNamespace(session=MagicMock())
|
||||
|
||||
def _fake_run(coro_or_factory, timeout=30):
|
||||
mcp._servers["playwright"] = connected
|
||||
coro = coro_or_factory() if callable(coro_or_factory) else coro_or_factory
|
||||
coro.close()
|
||||
return ["mcp_playwright_browser_navigate"]
|
||||
|
||||
with patch.object(mcp, "_ensure_mcp_loop"), \
|
||||
patch.object(mcp, "_run_on_mcp_loop", side_effect=_fake_run):
|
||||
assert mcp._ensure_lazy_server_connected("playwright") is True
|
||||
|
||||
assert "playwright" not in mcp._lazy_server_configs
|
||||
assert "playwright" not in mcp._lazy_server_fingerprints
|
||||
assert "playwright" not in mcp._lazy_server_tool_names
|
||||
|
||||
def test_lazy_connect_failure_records_cooldown(self):
|
||||
mcp._lazy_server_configs["playwright"] = {"command": "npx", "lazy": True}
|
||||
|
||||
def _fake_run(coro_or_factory, timeout=30):
|
||||
coro = coro_or_factory() if callable(coro_or_factory) else coro_or_factory
|
||||
coro.close()
|
||||
raise RuntimeError("spawn failed")
|
||||
|
||||
with patch.object(mcp, "_ensure_mcp_loop"), \
|
||||
patch.object(mcp, "_run_on_mcp_loop", side_effect=_fake_run), \
|
||||
patch.object(mcp, "_record_connect_failure") as mock_record:
|
||||
assert mcp._ensure_lazy_server_connected("playwright") is False
|
||||
|
||||
mock_record.assert_called_once_with("playwright")
|
||||
# Config retained so a later call can retry after cooldown.
|
||||
assert "playwright" in mcp._lazy_server_configs
|
||||
|
||||
|
||||
class TestResolveServerLazy:
|
||||
def test_default_off(self):
|
||||
assert mcp._resolve_server_lazy("s", {"command": "npx"}) is False
|
||||
|
||||
def test_explicit_true(self):
|
||||
assert mcp._resolve_server_lazy("s", {"command": "npx", "lazy": True}) is True
|
||||
|
||||
def test_explicit_false(self):
|
||||
assert mcp._resolve_server_lazy("s", {"command": "npx", "lazy": False}) is False
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
"""Unit tests for the on-disk MCP schema cache (tools/mcp_schema_cache.py).
|
||||
|
||||
The module landed in #56832's extraction without its tests; these cover the
|
||||
fingerprint keying, read/write round-trip, and invalidation behavior.
|
||||
"""
|
||||
|
||||
import tools.mcp_schema_cache as msc
|
||||
|
||||
|
||||
class TestConfigFingerprint:
|
||||
def test_stable_for_same_config(self):
|
||||
cfg = {"command": "npx", "args": ["-y", "@playwright/mcp"]}
|
||||
assert msc.config_fingerprint(cfg) == msc.config_fingerprint(dict(cfg))
|
||||
|
||||
def test_changes_when_connection_config_changes(self):
|
||||
base = {"command": "npx", "args": ["-y", "@playwright/mcp"]}
|
||||
assert msc.config_fingerprint(base) != msc.config_fingerprint(
|
||||
{**base, "args": ["-y", "@playwright/mcp", "--headless"]}
|
||||
)
|
||||
assert msc.config_fingerprint(base) != msc.config_fingerprint(
|
||||
{**base, "command": "uvx"}
|
||||
)
|
||||
assert msc.config_fingerprint(base) != msc.config_fingerprint(
|
||||
{**base, "tools": {"include": ["a"]}}
|
||||
)
|
||||
|
||||
def test_ignores_non_connection_keys(self):
|
||||
base = {"command": "npx", "args": []}
|
||||
assert msc.config_fingerprint(base) == msc.config_fingerprint(
|
||||
{**base, "timeout": 5, "enabled": True, "lazy": True}
|
||||
)
|
||||
|
||||
|
||||
class TestCacheRoundTrip:
|
||||
def _isolate(self, monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(msc, "_cache_path", lambda: tmp_path / "cache.json")
|
||||
|
||||
def test_write_then_read_with_matching_fingerprint(self, monkeypatch, tmp_path):
|
||||
self._isolate(monkeypatch, tmp_path)
|
||||
tools = [{"name": "t1", "description": "d", "inputSchema": {"type": "object"}}]
|
||||
msc.write_cache_entry("srv", "fp1", tools=tools, utility_tools=[])
|
||||
entry = msc.get_cached_entry("srv", "fp1")
|
||||
assert entry is not None
|
||||
assert msc.tools_from_cache_entry(entry) == tools
|
||||
assert msc.utility_tools_from_cache_entry(entry) == []
|
||||
assert msc.has_cached_entry("srv", "fp1")
|
||||
|
||||
def test_fingerprint_mismatch_returns_none(self, monkeypatch, tmp_path):
|
||||
self._isolate(monkeypatch, tmp_path)
|
||||
msc.write_cache_entry("srv", "fp1", tools=[], utility_tools=[])
|
||||
assert msc.get_cached_entry("srv", "OTHER") is None
|
||||
assert not msc.has_cached_entry("srv", "OTHER")
|
||||
|
||||
def test_missing_server_returns_none(self, monkeypatch, tmp_path):
|
||||
self._isolate(monkeypatch, tmp_path)
|
||||
assert msc.get_cached_entry("nope", "fp") is None
|
||||
|
||||
def test_clear_cache_entry(self, monkeypatch, tmp_path):
|
||||
self._isolate(monkeypatch, tmp_path)
|
||||
msc.write_cache_entry("srv", "fp1", tools=[], utility_tools=[])
|
||||
msc.clear_cache_entry("srv")
|
||||
assert msc.get_cached_entry("srv", "fp1") is None
|
||||
|
||||
def test_corrupt_cache_file_is_tolerated(self, monkeypatch, tmp_path):
|
||||
self._isolate(monkeypatch, tmp_path)
|
||||
(tmp_path / "cache.json").write_text("{not json", encoding="utf-8")
|
||||
assert msc.get_cached_entry("srv", "fp") is None
|
||||
# And writes recover the file.
|
||||
msc.write_cache_entry("srv", "fp", tools=[], utility_tools=[])
|
||||
assert msc.has_cached_entry("srv", "fp")
|
||||
|
||||
def test_malformed_entry_shapes_are_tolerated(self):
|
||||
assert msc.tools_from_cache_entry({"tools": "nope"}) == []
|
||||
assert msc.utility_tools_from_cache_entry({}) == []
|
||||
|
|
@ -3532,6 +3532,12 @@ class MCPServerTask:
|
|||
_servers: Dict[str, MCPServerTask] = {}
|
||||
_server_connecting: set[str] = set()
|
||||
_server_connect_errors: Dict[str, str] = {}
|
||||
# Lazy MCP startup (#56832): servers whose tools were registered from the
|
||||
# on-disk schema cache without spawning/connecting. Keyed by server name;
|
||||
# entries are popped once a real connection is established on first use.
|
||||
_lazy_server_configs: Dict[str, dict] = {}
|
||||
_lazy_server_fingerprints: Dict[str, str] = {}
|
||||
_lazy_server_tool_names: Dict[str, List[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.
|
||||
|
|
@ -4758,10 +4764,84 @@ def _request_lazy_reconnect(server_name: str, server: MCPServerTask) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _get_connected_server_for_call(server_name: str) -> Optional[MCPServerTask]:
|
||||
"""Return a connected server, lazily reconnecting recycled stdio state."""
|
||||
def _resolve_server_lazy(name: str, config: dict) -> bool:
|
||||
"""True when this server defers spawn/connect until first tool use.
|
||||
|
||||
Gated per-server by ``mcp_servers.<name>.lazy`` in config (default OFF),
|
||||
following the same per-server key pattern as ``idle_timeout_seconds``.
|
||||
Design from #56832 (Vansh5632).
|
||||
"""
|
||||
return _parse_boolish(config.get("lazy", False), default=False)
|
||||
|
||||
|
||||
def _ensure_lazy_server_connected(server_name: str) -> bool:
|
||||
"""Connect a lazily-registered MCP server on demand (sync, blocks caller).
|
||||
|
||||
Composes with the existing connect machinery: respects the per-server
|
||||
connect cooldown (#50394), the ``_server_connecting`` dedup set, and
|
||||
routes through ``_discover_and_register_server`` so parked/recycle/
|
||||
cooldown bookkeeping stays in one place. Returns True when a live
|
||||
session is available afterwards.
|
||||
"""
|
||||
with _lock:
|
||||
server = _servers.get(server_name)
|
||||
if server is not None and server.session is not None:
|
||||
return True
|
||||
config = _lazy_server_configs.get(server_name)
|
||||
if not config:
|
||||
return False
|
||||
if _connect_cooldown_active(server_name):
|
||||
return False
|
||||
if server_name in _server_connecting:
|
||||
return False
|
||||
_server_connecting.add(server_name)
|
||||
_server_connect_errors.pop(server_name, None)
|
||||
|
||||
logger.info("MCP server '%s': lazy start on first use", server_name)
|
||||
_ensure_mcp_loop()
|
||||
connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT)
|
||||
|
||||
async def _connect():
|
||||
return await _discover_and_register_server(server_name, config)
|
||||
|
||||
try:
|
||||
_run_on_mcp_loop(_connect, timeout=float(connect_timeout) + 30.0)
|
||||
except BaseException as exc:
|
||||
message = _format_connect_error(exc)
|
||||
with _lock:
|
||||
_server_connecting.discard(server_name)
|
||||
_server_connect_errors[server_name] = message
|
||||
_record_connect_failure(server_name)
|
||||
logger.warning(
|
||||
"Lazy MCP connect failed for '%s': %s", server_name, message,
|
||||
)
|
||||
return False
|
||||
|
||||
with _lock:
|
||||
_server_connecting.discard(server_name)
|
||||
_clear_connect_failure(server_name)
|
||||
_lazy_server_configs.pop(server_name, None)
|
||||
_lazy_server_fingerprints.pop(server_name, None)
|
||||
_lazy_server_tool_names.pop(server_name, None)
|
||||
server = _servers.get(server_name)
|
||||
return server is not None and server.session is not None
|
||||
|
||||
|
||||
def _get_connected_server_for_call(server_name: str) -> Optional[MCPServerTask]:
|
||||
"""Return a connected server, lazily reconnecting recycled stdio state.
|
||||
|
||||
Also the single first-use connect point for lazy (schema-cache
|
||||
registered) servers, so raw tool calls AND the resource/prompt utility
|
||||
handlers all trigger the deferred spawn (#56832).
|
||||
"""
|
||||
with _lock:
|
||||
server = _servers.get(server_name)
|
||||
is_lazy = server_name in _lazy_server_configs
|
||||
if is_lazy and (server is None or server.session is None):
|
||||
_ensure_lazy_server_connected(server_name)
|
||||
with _lock:
|
||||
server = _servers.get(server_name)
|
||||
return server
|
||||
if server is not None and server.session is None and server._is_recycled_stdio():
|
||||
_request_lazy_reconnect(server_name, server)
|
||||
with _lock:
|
||||
|
|
@ -5240,10 +5320,13 @@ def _make_check_fn(server_name: str):
|
|||
def _check() -> bool:
|
||||
with _lock:
|
||||
server = _servers.get(server_name)
|
||||
return (
|
||||
server is not None
|
||||
and (server.session is not None or server._is_recycled_stdio())
|
||||
)
|
||||
if server is not None and (
|
||||
server.session is not None or server._is_recycled_stdio()
|
||||
):
|
||||
return True
|
||||
# Lazy (schema-cache registered) servers are available: the
|
||||
# first real call spawns/connects them (#56832).
|
||||
return server_name in _lazy_server_configs
|
||||
|
||||
return _check
|
||||
|
||||
|
|
@ -5692,6 +5775,16 @@ def _existing_tool_names() -> List[str]:
|
|||
for mcp_tool in server._tools:
|
||||
schema = _convert_mcp_schema(server.name, mcp_tool)
|
||||
names.append(schema["name"])
|
||||
# Lazy servers registered from the schema cache have no MCPServerTask
|
||||
# yet — their tools live in the registry only (#56832).
|
||||
with _lock:
|
||||
lazy_names = [
|
||||
n
|
||||
for sname, tool_names in _lazy_server_tool_names.items()
|
||||
if sname not in _servers
|
||||
for n in tool_names
|
||||
]
|
||||
names.extend(lazy_names)
|
||||
return names
|
||||
|
||||
|
||||
|
|
@ -5879,7 +5972,162 @@ def _register_server_tools(name: str, server: MCPServerTask, config: dict) -> Li
|
|||
|
||||
if registered_names:
|
||||
registry.register_toolset_alias(name, toolset_name)
|
||||
# Write-through (#56832): refresh the on-disk schema cache after a
|
||||
# live connect so the next startup can lazily register this server
|
||||
# without spawning it. Cache failures never break registration.
|
||||
try:
|
||||
from tools.mcp_schema_cache import config_fingerprint, write_cache_entry
|
||||
|
||||
tools_payload: List[dict] = []
|
||||
for mcp_tool in server._tools:
|
||||
if not _should_register(mcp_tool.name):
|
||||
continue
|
||||
schema_obj = getattr(mcp_tool, "inputSchema", None)
|
||||
tools_payload.append({
|
||||
"name": mcp_tool.name,
|
||||
"description": mcp_tool.description or "",
|
||||
"inputSchema": schema_obj if isinstance(schema_obj, dict) else {},
|
||||
})
|
||||
utility_payload = [
|
||||
{"schema": entry["schema"], "handler_key": entry["handler_key"]}
|
||||
for entry in _select_utility_schemas(name, server, config)
|
||||
]
|
||||
write_cache_entry(
|
||||
name,
|
||||
config_fingerprint(config),
|
||||
tools=tools_payload,
|
||||
utility_tools=utility_payload,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.debug("MCP schema cache write failed for '%s': %s", name, exc)
|
||||
|
||||
return registered_names
|
||||
|
||||
|
||||
class _CachedMCPTool:
|
||||
"""Minimal stand-in for MCP Tool objects loaded from the schema cache."""
|
||||
|
||||
__slots__ = ("name", "description", "inputSchema")
|
||||
|
||||
def __init__(self, name: str, description: str, inputSchema: dict):
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.inputSchema = inputSchema or {}
|
||||
|
||||
|
||||
def _register_from_cache_sync(name: str, config: dict, entry: dict) -> List[str]:
|
||||
"""Register a server's tools from a cached manifest, no child process.
|
||||
|
||||
Lazy startup (#56832, design by Vansh5632): tools appear in the registry
|
||||
immediately; the first real call routes through
|
||||
``_get_connected_server_for_call`` → ``_ensure_lazy_server_connected``.
|
||||
"""
|
||||
from tools.registry import registry
|
||||
from tools.mcp_schema_cache import (
|
||||
config_fingerprint,
|
||||
tools_from_cache_entry,
|
||||
utility_tools_from_cache_entry,
|
||||
)
|
||||
|
||||
registered_names: List[str] = []
|
||||
toolset_name = f"mcp-{name}"
|
||||
fingerprint = config_fingerprint(config)
|
||||
tool_timeout = config.get("timeout", _DEFAULT_TOOL_TIMEOUT)
|
||||
tools_filter = config.get("tools") or {}
|
||||
include_set = _normalize_name_filter(
|
||||
tools_filter.get("include"), f"mcp_servers.{name}.tools.include"
|
||||
)
|
||||
exclude_set = _normalize_name_filter(
|
||||
tools_filter.get("exclude"), f"mcp_servers.{name}.tools.exclude"
|
||||
)
|
||||
|
||||
def _should_register(tool_name: str) -> bool:
|
||||
if include_set:
|
||||
return matches_name_filter(tool_name, include_set)
|
||||
if exclude_set:
|
||||
return not matches_name_filter(tool_name, exclude_set)
|
||||
return True
|
||||
|
||||
check_fn = _make_check_fn(name)
|
||||
for raw in tools_from_cache_entry(entry):
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
raw_name = raw.get("name")
|
||||
if not raw_name or not _should_register(raw_name):
|
||||
continue
|
||||
raw_schema = raw.get("inputSchema")
|
||||
mcp_tool = _CachedMCPTool(
|
||||
raw_name,
|
||||
raw.get("description") or "",
|
||||
raw_schema if isinstance(raw_schema, dict) else {},
|
||||
)
|
||||
schema = _convert_mcp_schema(name, mcp_tool)
|
||||
registry_name = schema["name"]
|
||||
existing_toolset = registry.get_toolset_for_tool(registry_name)
|
||||
if existing_toolset and existing_toolset != toolset_name:
|
||||
logger.warning(
|
||||
"MCP server '%s' (lazy): cached tool '%s' collides with "
|
||||
"toolset '%s' — skipping",
|
||||
name, registry_name, existing_toolset,
|
||||
)
|
||||
continue
|
||||
registry.register(
|
||||
name=registry_name,
|
||||
toolset=toolset_name,
|
||||
schema=schema,
|
||||
handler=_make_tool_handler(name, raw_name, tool_timeout),
|
||||
check_fn=check_fn,
|
||||
is_async=False,
|
||||
description=schema["description"],
|
||||
)
|
||||
if registry.get_toolset_for_tool(registry_name) != toolset_name:
|
||||
continue
|
||||
_track_mcp_tool_server(registry_name, name)
|
||||
registered_names.append(registry_name)
|
||||
|
||||
handler_factories = {
|
||||
"list_resources": _make_list_resources_handler,
|
||||
"read_resource": _make_read_resource_handler,
|
||||
"list_prompts": _make_list_prompts_handler,
|
||||
"get_prompt": _make_get_prompt_handler,
|
||||
}
|
||||
for raw in utility_tools_from_cache_entry(entry):
|
||||
if not isinstance(raw, dict):
|
||||
continue
|
||||
schema = raw.get("schema")
|
||||
handler_key = raw.get("handler_key")
|
||||
if not isinstance(schema, dict) or handler_key not in handler_factories:
|
||||
continue
|
||||
util_name = schema.get("name") or ""
|
||||
if not util_name:
|
||||
continue
|
||||
existing_toolset = registry.get_toolset_for_tool(util_name)
|
||||
if existing_toolset and existing_toolset != toolset_name:
|
||||
continue
|
||||
registry.register(
|
||||
name=util_name,
|
||||
toolset=toolset_name,
|
||||
schema=schema,
|
||||
handler=handler_factories[handler_key](name, tool_timeout),
|
||||
check_fn=check_fn,
|
||||
is_async=False,
|
||||
description=schema.get("description") or "",
|
||||
)
|
||||
if registry.get_toolset_for_tool(util_name) != toolset_name:
|
||||
continue
|
||||
_track_mcp_tool_server(util_name, name)
|
||||
registered_names.append(util_name)
|
||||
|
||||
if registered_names:
|
||||
registry.register_toolset_alias(name, toolset_name)
|
||||
with _lock:
|
||||
_lazy_server_configs[name] = dict(config)
|
||||
_lazy_server_fingerprints[name] = fingerprint
|
||||
_lazy_server_tool_names[name] = list(registered_names)
|
||||
logger.info(
|
||||
"MCP server '%s' (lazy): registered %d tool(s) from schema cache",
|
||||
name, len(registered_names),
|
||||
)
|
||||
return registered_names
|
||||
|
||||
async def _discover_and_register_server(name: str, config: dict) -> List[str]:
|
||||
|
|
@ -5980,6 +6228,9 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
|
|||
for k, v in servers.items()
|
||||
if k not in _servers
|
||||
and k not in connecting
|
||||
# Servers already lazily registered from the schema cache are
|
||||
# not re-registered; they connect on first tool use (#56832).
|
||||
and k not in _lazy_server_configs
|
||||
and _parse_boolish(v.get("enabled", True), default=True)
|
||||
# Skip a server still serving its post-failure backoff. Without
|
||||
# this, a server that fails to connect (and is therefore never
|
||||
|
|
@ -6015,6 +6266,51 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
|
|||
if not new_servers:
|
||||
return _existing_tool_names()
|
||||
|
||||
# Lazy startup (#56832): servers gated with ``lazy: true`` whose config
|
||||
# fingerprint matches a valid on-disk schema-cache entry register their
|
||||
# tools from cache WITHOUT spawning/connecting. A missing or stale cache
|
||||
# entry falls back to the normal eager connect below (which write-through
|
||||
# refreshes the cache for next time).
|
||||
eager_servers: Dict[str, dict] = dict(new_servers)
|
||||
lazy_registered = 0
|
||||
lazy_server_count = 0
|
||||
try:
|
||||
from tools.mcp_schema_cache import config_fingerprint, get_cached_entry
|
||||
except Exception: # pragma: no cover - cache module missing
|
||||
config_fingerprint = None # type: ignore[assignment]
|
||||
get_cached_entry = None # type: ignore[assignment]
|
||||
if config_fingerprint is not None and get_cached_entry is not None:
|
||||
for name, cfg in new_servers.items():
|
||||
if not _resolve_server_lazy(name, cfg):
|
||||
continue
|
||||
entry = get_cached_entry(name, config_fingerprint(cfg))
|
||||
if not entry:
|
||||
continue
|
||||
with _lock:
|
||||
_server_connecting.discard(name)
|
||||
try:
|
||||
names = _register_from_cache_sync(name, cfg, entry)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed lazy MCP registration for '%s': %s", name, exc,
|
||||
)
|
||||
with _lock:
|
||||
_server_connecting.add(name)
|
||||
continue
|
||||
eager_servers.pop(name, None)
|
||||
lazy_registered += len(names)
|
||||
lazy_server_count += 1
|
||||
new_servers = eager_servers
|
||||
|
||||
if not new_servers:
|
||||
if lazy_registered:
|
||||
logger.info(
|
||||
"MCP: registered %d lazy tool(s) from schema cache "
|
||||
"(no processes spawned)",
|
||||
lazy_registered,
|
||||
)
|
||||
return _existing_tool_names()
|
||||
|
||||
# Start the background event loop for MCP connections
|
||||
_ensure_mcp_loop()
|
||||
|
||||
|
|
@ -6103,8 +6399,10 @@ def register_mcp_servers(servers: Dict[str, dict]) -> List[str]:
|
|||
for n in connected
|
||||
)
|
||||
failed = len(new_servers) - len(connected)
|
||||
new_tool_count += lazy_registered
|
||||
connected_count = len(connected) + lazy_server_count
|
||||
if new_tool_count or failed:
|
||||
summary = f"MCP: registered {new_tool_count} tool(s) from {len(connected)} server(s)"
|
||||
summary = f"MCP: registered {new_tool_count} tool(s) from {connected_count} server(s)"
|
||||
if failed:
|
||||
summary += f" ({failed} failed)"
|
||||
logger.info(summary)
|
||||
|
|
|
|||
Loading…
Reference in New Issue