diff --git a/tests/tools/test_tool_manifest_cache.py b/tests/tools/test_tool_manifest_cache.py new file mode 100644 index 00000000000..6ddd118bb44 --- /dev/null +++ b/tests/tools/test_tool_manifest_cache.py @@ -0,0 +1,153 @@ +"""Tests for the mtime-based tool discovery manifest cache. + +The manifest cache avoids re-parsing 31 Python files with AST on every +process start when none of them have changed since the last run. +""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +import time +from pathlib import Path + +import pytest + + +def _load_registry_module(): + """Import tools/registry.py fresh so module-level state is clean.""" + spec = importlib.util.spec_from_file_location( + "_registry_under_test", + Path(__file__).resolve().parents[2] / "tools" / "registry.py", + ) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture() +def tool_dir(tmp_path): + """Create a minimal tools directory with two self-registering files.""" + tools = tmp_path / "tools" + tools.mkdir() + (tools / "__init__.py").write_text("", encoding="utf-8") + (tools / "registry.py").write_text("# stub\n", encoding="utf-8") + (tools / "alpha.py").write_text( + "from tools.registry import registry\n" + "registry.register(name='alpha', toolset='a', schema={}, handler=lambda *_a, **_k: '{}')\n", + encoding="utf-8", + ) + (tools / "beta.py").write_text( + "from tools.registry import registry\n" + "registry.register(name='beta', toolset='b', schema={}, handler=lambda *_a, **_k: '{}')\n", + encoding="utf-8", + ) + (tools / "helper.py").write_text("VALUE = 1\n", encoding="utf-8") + return tools + + +@pytest.fixture() +def cache_path(tmp_path): + return tmp_path / "cache" / "tool_manifest.json" + + +class TestManifestCache: + def test_save_and_load_roundtrip(self, tool_dir, cache_path): + module = _load_registry_module() + module._save_manifest_cache(cache_path, tools_dir=tool_dir, module_names=["tools.alpha", "tools.beta"]) + assert cache_path.exists() + loaded = module._load_manifest_cache(cache_path, tools_dir=tool_dir) + assert loaded is not None + assert loaded == ["tools.alpha", "tools.beta"] + + def test_cache_miss_when_file_missing(self, tool_dir, cache_path): + module = _load_registry_module() + assert module._load_manifest_cache(cache_path, tools_dir=tool_dir) is None + + def test_cache_miss_when_mtime_changes(self, tool_dir, cache_path): + module = _load_registry_module() + module._save_manifest_cache(cache_path, tools_dir=tool_dir, module_names=["tools.alpha", "tools.beta"]) + time.sleep(0.05) + (tool_dir / "alpha.py").touch() + assert module._load_manifest_cache(cache_path, tools_dir=tool_dir) is None + + def test_cache_miss_when_file_added(self, tool_dir, cache_path): + module = _load_registry_module() + module._save_manifest_cache(cache_path, tools_dir=tool_dir, module_names=["tools.alpha", "tools.beta"]) + (tool_dir / "gamma.py").write_text( + "from tools.registry import registry\n" + "registry.register(name='gamma', toolset='g', schema={}, handler=lambda *_a, **_k: '{}')\n", + encoding="utf-8", + ) + assert module._load_manifest_cache(cache_path, tools_dir=tool_dir) is None + + def test_cache_miss_when_file_removed(self, tool_dir, cache_path): + module = _load_registry_module() + module._save_manifest_cache(cache_path, tools_dir=tool_dir, module_names=["tools.alpha", "tools.beta"]) + (tool_dir / "beta.py").unlink() + assert module._load_manifest_cache(cache_path, tools_dir=tool_dir) is None + + def test_cache_miss_when_corrupted_json(self, tool_dir, cache_path): + module = _load_registry_module() + cache_path.parent.mkdir(parents=True, exist_ok=True) + cache_path.write_text("not json", encoding="utf-8") + assert module._load_manifest_cache(cache_path, tools_dir=tool_dir) is None + + def test_discover_uses_cache_to_skip_ast_scan(self, tool_dir, cache_path, monkeypatch): + import unittest.mock + + module = _load_registry_module() + module._save_manifest_cache(cache_path, tools_dir=tool_dir, module_names=["tools.alpha", "tools.beta"]) + ast_calls = {"count": 0} + original_scan = module._module_registers_tools + + def counting_scan(path): + ast_calls["count"] += 1 + return original_scan(path) + + monkeypatch.setattr(module, "_module_registers_tools", counting_scan) + with unittest.mock.patch.object(module.importlib, "import_module", side_effect=lambda name: None): + result = module.discover_builtin_tools(tools_dir=tool_dir, manifest_cache_path=cache_path) + assert sorted(result) == ["tools.alpha", "tools.beta"] + assert ast_calls["count"] == 0, "AST scan should be skipped when cache is valid" + + def test_discover_falls_back_to_scan_when_cache_invalid(self, tool_dir, cache_path, monkeypatch): + import unittest.mock + + module = _load_registry_module() + ast_calls = {"count": 0} + original_scan = module._module_registers_tools + + def counting_scan(path): + ast_calls["count"] += 1 + return original_scan(path) + + monkeypatch.setattr(module, "_module_registers_tools", counting_scan) + with unittest.mock.patch.object(module.importlib, "import_module", side_effect=lambda name: None): + result = module.discover_builtin_tools(tools_dir=tool_dir, manifest_cache_path=cache_path) + assert sorted(result) == ["tools.alpha", "tools.beta"] + assert ast_calls["count"] > 0, "AST scan should run when cache is invalid" + assert cache_path.exists(), "Cache should be written after fallback scan" + + def test_discover_skips_cache_when_env_disabled(self, tool_dir, cache_path, monkeypatch): + import unittest.mock + + module = _load_registry_module() + module._save_manifest_cache(cache_path, tools_dir=tool_dir, module_names=["tools.alpha", "tools.beta"]) + monkeypatch.setenv("HERMES_NO_TOOL_CACHE", "1") + ast_calls = {"count": 0} + original_scan = module._module_registers_tools + + def counting_scan(path): + ast_calls["count"] += 1 + return original_scan(path) + + monkeypatch.setattr(module, "_module_registers_tools", counting_scan) + with unittest.mock.patch.object(module.importlib, "import_module", side_effect=lambda name: None): + result = module.discover_builtin_tools(tools_dir=tool_dir, manifest_cache_path=cache_path) + assert sorted(result) == ["tools.alpha", "tools.beta"] + assert ast_calls["count"] > 0, "AST scan should run when cache is disabled via env" diff --git a/tools/registry.py b/tools/registry.py index ede7c487dd9..8eb18bc6e00 100644 --- a/tools/registry.py +++ b/tools/registry.py @@ -18,6 +18,7 @@ import ast import importlib import json import logging +import os import sys import threading import time @@ -64,15 +65,89 @@ def _module_registers_tools(module_path: Path) -> bool: return any(_is_registry_register_call(stmt) for stmt in tree.body) -def discover_builtin_tools(tools_dir: Optional[Path] = None) -> List[str]: - """Import built-in self-registering tool modules and return their module names.""" +def _default_manifest_cache_path() -> Optional[Path]: + """Return the default manifest cache path under HERMES_HOME, or None.""" + try: + from hermes_constants import get_hermes_home + hermes_home = get_hermes_home() + except Exception: + return None + return Path(hermes_home) / "cache" / "tool_manifest.json" + + +def _collect_tool_file_mtimes(tools_dir: Path) -> Dict[str, float]: + """Return ``{relative_name: mtime}`` for every candidate tool file.""" + mtimes: Dict[str, float] = {} + for path in sorted(tools_dir.glob("*.py")): + if path.name in {"__init__.py", "registry.py", "mcp_tool.py"}: + continue + try: + mtimes[path.name] = path.stat().st_mtime + except OSError: + pass + return mtimes + + +def _save_manifest_cache(cache_path: Path, tools_dir: Path, module_names: List[str]) -> None: + """Persist the AST-scan result so the next startup can skip it.""" + try: + cache_path.parent.mkdir(parents=True, exist_ok=True) + payload = { + "version": 1, + "tools_dir": str(tools_dir), + "module_names": module_names, + "mtimes": _collect_tool_file_mtimes(tools_dir), + } + cache_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") + except Exception: + logger.debug("Failed to save tool manifest cache", exc_info=True) + + +def _load_manifest_cache(cache_path: Path, tools_dir: Path) -> Optional[List[str]]: + """Return cached module names if the cache is valid, else ``None``.""" + if cache_path is None: + return None + try: + data = json.loads(cache_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError, UnicodeDecodeError): + return None + if not isinstance(data, dict) or data.get("version") != 1: + return None + cached_mtimes = data.get("mtimes", {}) + current_mtimes = _collect_tool_file_mtimes(tools_dir) + if cached_mtimes != current_mtimes: + return None + module_names = data.get("module_names") + if not isinstance(module_names, list): + return None + return module_names + + +def discover_builtin_tools(tools_dir: Optional[Path] = None, manifest_cache_path: Optional[Path] = None) -> List[str]: + """Import built-in self-registering tool modules and return their module names. + + When *manifest_cache_path* is set (and ``HERMES_NO_TOOL_CACHE`` is not + ``"1"``), the function checks the mtime-based manifest cache first. On a + hit the AST scan of every ``tools/*.py`` file is skipped entirely, saving + ~500 ms on cold startup. + """ tools_path = Path(tools_dir) if tools_dir is not None else Path(__file__).resolve().parent - module_names = [ - f"tools.{path.stem}" - for path in sorted(tools_path.glob("*.py")) - if path.name not in {"__init__.py", "registry.py", "mcp_tool.py"} - and _module_registers_tools(path) - ] + if manifest_cache_path is None: + manifest_cache_path = _default_manifest_cache_path() + + module_names: Optional[List[str]] = None + if not os.environ.get("HERMES_NO_TOOL_CACHE"): + module_names = _load_manifest_cache(manifest_cache_path, tools_path) + + if module_names is None: + module_names = [ + f"tools.{path.stem}" + for path in sorted(tools_path.glob("*.py")) + if path.name not in {"__init__.py", "registry.py", "mcp_tool.py"} + and _module_registers_tools(path) + ] + if manifest_cache_path is not None: + _save_manifest_cache(manifest_cache_path, tools_path, module_names) imported: List[str] = [] for mod_name in module_names: