diff --git a/api/streaming.py b/api/streaming.py index 7968c6ed..44d397ff 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -40,6 +40,30 @@ from api.metering import meter # save/restore around the entire agent run. _ENV_LOCK = threading.Lock() + +def _prewarm_skill_tool_modules(): + """Import tools.skills_tool and tools.skill_manager_tool outside any lock. + + First-time module imports can trigger heavy initialisation (disk I/O, + transitive imports, plugin discovery). Performing those imports while + holding ``_ENV_LOCK`` serialises every concurrent session behind the + slowest import. Prewarming ensures the modules are already in + ``sys.modules`` before the lock is acquired, so the lock body only + does lightweight attribute patching. + + We cannot place these at module top-level because ``tools.*`` lives + in the hermes-agent package which may not be on ``sys.path`` at + import time (Docker volume-mount ordering). A dedicated helper + keeps the lazy-import try/except in one place and makes the intent + explicit. + """ + for _mod_name in ('tools.skills_tool', 'tools.skill_manager_tool'): + try: + __import__(_mod_name) + except ImportError: + pass + + # Lazy import to avoid circular deps -- hermes-agent is on sys.path via api/config.py try: from run_agent import AIAgent @@ -2240,6 +2264,10 @@ def _run_agent_streaming( _profile_home, ) _set_thread_env(**_thread_env) + # Prewarm skill-tool imports *before* acquiring the lock so that + # first-time module initialisation (which can be slow) does not + # block other concurrent sessions waiting on _ENV_LOCK (#2024). + _prewarm_skill_tool_modules() # Still set process-level env as fallback for tools that bypass thread-local # Acquire lock only for the env mutation, then release before the agent runs. # The finally block re-acquires to restore — keeping critical sections short @@ -2259,20 +2287,27 @@ def _run_agent_streaming( # Patch module-level caches to match the active profile. # _set_hermes_home() does this for process-wide switches # but per-request switches skip it (#1700). + # Modules were prewarmed by _prewarm_skill_tool_modules() + # above, so we only do lightweight sys.modules lookups and + # attribute assignments here — no first-time import under + # the lock (#2024). from pathlib import Path as _P + import sys as _sys _ph = _P(_profile_home) - try: - import tools.skills_tool as _sk - _sk.HERMES_HOME = _ph - _sk.SKILLS_DIR = _ph / 'skills' - except (ImportError, AttributeError): - pass - try: - import tools.skill_manager_tool as _sm - _sm.HERMES_HOME = _ph - _sm.SKILLS_DIR = _ph / 'skills' - except (ImportError, AttributeError): - pass + _sk = _sys.modules.get('tools.skills_tool') + if _sk is not None: + try: + _sk.HERMES_HOME = _ph + _sk.SKILLS_DIR = _ph / 'skills' + except AttributeError: + pass + _sm = _sys.modules.get('tools.skill_manager_tool') + if _sm is not None: + try: + _sm.HERMES_HOME = _ph + _sm.SKILLS_DIR = _ph / 'skills' + except AttributeError: + pass # Lock released — agent runs without holding it # ── MCP Server Discovery (lazy import, idempotent) ── # MUST run AFTER the HERMES_HOME mutation above — `discover_mcp_tools()` diff --git a/tests/test_issue2024_env_lock_skill_imports.py b/tests/test_issue2024_env_lock_skill_imports.py new file mode 100644 index 00000000..66250b91 --- /dev/null +++ b/tests/test_issue2024_env_lock_skill_imports.py @@ -0,0 +1,219 @@ +"""Regression test for issue #2024. + +tools.skills_tool / tools.skill_manager_tool imports must NOT appear +inside an ``_ENV_LOCK`` body in api/streaming.py. First-time module +imports can be slow (disk I/O, transitive deps, plugin discovery) and +holding the lock during them serialises every concurrent session behind +the slowest import. + +The fix introduces ``_prewarm_skill_tool_modules()`` which does the +imports *before* the lock is acquired, and the lock body uses only +``sys.modules.get()`` lookups (O(1) dict lookup, no import machinery). + +These tests are AST/source-level because the actual import targets +(``tools.skills_tool``, ``tools.skill_manager_tool``) live in the +hermes-agent package which may not be installed in the test venv. +""" +import ast +import pathlib +import textwrap + +REPO = pathlib.Path(__file__).resolve().parent.parent +STREAMING_PY = REPO / "api" / "streaming.py" + + +def _read_streaming() -> str: + return STREAMING_PY.read_text(encoding="utf-8") + + +# --------------------------------------------------------------------------- +# AST-level check: walk every ``with`` statement whose context-expression +# references ``_ENV_LOCK`` and ensure no ``Import`` or ``ImportFrom`` +# node for the two target modules exists in its body. +# --------------------------------------------------------------------------- + +def _find_env_lock_with_bodies(source: str) -> list[list[ast.stmt]]: + """Return the statement-list bodies of all ``with _ENV_LOCK:`` blocks.""" + tree = ast.parse(source) + bodies: list[list[ast.stmt]] = [] + + class _Visitor(ast.NodeVisitor): + def visit_With(self, node: ast.With): + # Check whether any context-expression is a simple Name `_ENV_LOCK` + for item in node.items: + ctx = item.context_expr + if isinstance(ctx, ast.Name) and ctx.id == "_ENV_LOCK": + bodies.append(node.body) + break + self.generic_visit(node) + + _Visitor().visit(tree) + return bodies + + +def _imports_in_body(body: list[ast.stmt], target_modules: set[str]) -> list[str]: + """Return module names from Import/ImportFrom nodes in *body* that are in *target_modules*.""" + found: list[str] = [] + for node in ast.walk(ast.Module(body=body, type_ignores=[])): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name in target_modules: + found.append(alias.name) + elif isinstance(node, ast.ImportFrom): + if node.module in target_modules: + found.append(node.module) + return found + + +_TARGET_MODULES = {"tools.skills_tool", "tools.skill_manager_tool"} + + +class TestNoSkillToolImportsInsideEnvLock: + """AST-level: no ``import tools.skills_tool`` or ``import tools.skill_manager_tool`` + inside any ``with _ENV_LOCK:`` block.""" + + def test_no_skill_imports_in_env_lock(self): + source = _read_streaming() + bodies = _find_env_lock_with_bodies(source) + assert bodies, "Expected at least one `with _ENV_LOCK:` block in streaming.py" + for body in bodies: + found = _imports_in_body(body, _TARGET_MODULES) + assert found == [], ( + f"Found import(s) of {found} inside an `_ENV_LOCK` with-block. " + "Move them to _prewarm_skill_tool_modules() outside the lock (#2024)." + ) + + +class TestPrewarmHelperExists: + """The ``_prewarm_skill_tool_modules`` helper must exist and reference + both target modules.""" + + def test_prewarm_function_defined(self): + source = _read_streaming() + tree = ast.parse(source) + func_names = { + node.name + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) + } + assert "_prewarm_skill_tool_modules" in func_names, ( + "_prewarm_skill_tool_modules() helper must be defined in streaming.py" + ) + + def test_prewarm_references_both_modules(self): + source = _read_streaming() + # Find the function source and check it references both module names. + # Simple string check is sufficient and more robust than AST for + # dynamic __import__ calls. + assert "tools.skills_tool" in source, ( + "streaming.py must reference 'tools.skills_tool'" + ) + assert "tools.skill_manager_tool" in source, ( + "streaming.py must reference 'tools.skill_manager_tool'" + ) + + def test_prewarm_called_before_env_lock(self): + """_prewarm_skill_tool_modules() must be called before the first + ``with _ENV_LOCK:`` in _run_agent_streaming.""" + source = _read_streaming() + lines = source.splitlines() + prewarm_line = None + first_env_lock_line = None + for i, line in enumerate(lines, 1): + if "_prewarm_skill_tool_modules()" in line and prewarm_line is None: + prewarm_line = i + if "with _ENV_LOCK:" in line and first_env_lock_line is None: + first_env_lock_line = i + assert prewarm_line is not None, "_prewarm_skill_tool_modules() call not found" + assert first_env_lock_line is not None, "with _ENV_LOCK: not found" + assert prewarm_line < first_env_lock_line, ( + f"_prewarm_skill_tool_modules() (line {prewarm_line}) must appear " + f"before the first `with _ENV_LOCK:` (line {first_env_lock_line})" + ) + + +class TestSysModulesLookupInEnvLock: + """Inside the lock, the code must use ``sys.modules.get()`` instead of + ``import`` for the skill-tool modules.""" + + def test_sys_modules_get_used_in_env_lock(self): + source = _read_streaming() + bodies = _find_env_lock_with_bodies(source) + assert bodies, "Expected at least one `with _ENV_LOCK:` block" + + # Collect all string content within the lock bodies by extracting + # Constant/Str nodes — simpler than full AST string reconstruction. + lock_source_segments: list[str] = [] + for body in bodies: + for node in ast.walk(ast.Module(body=body, type_ignores=[])): + if isinstance(node, ast.Constant) and isinstance(node.value, str): + lock_source_segments.append(node.value) + + # The lock body should reference sys.modules.get for both modules + lock_text = "\n".join(lock_source_segments) + # More reliable: check the raw source lines inside the lock + lines = source.splitlines() + in_lock = False + lock_lines: list[str] = [] + depth = 0 + for line in lines: + stripped = line.strip() + if stripped.startswith("with _ENV_LOCK:"): + in_lock = True + depth = 0 + continue + if in_lock: + # Track indentation depth to know when we exit the with-block + if stripped: + # Count leading spaces + indent = len(line) - len(line.lstrip()) + if depth == 0: + depth = indent + elif indent < depth and stripped: + in_lock = False + continue + lock_lines.append(line) + + lock_source = "\n".join(lock_lines) + assert "sys.modules.get" in lock_source, ( + "Inside `_ENV_LOCK`, skill-tool modules must be accessed via " + "`sys.modules.get()` instead of `import` (#2024)" + ) + assert "tools.skills_tool" in lock_source, ( + "tools.skills_tool must still be referenced inside `_ENV_LOCK` " + "for attribute patching (HERMES_HOME / SKILLS_DIR)" + ) + assert "tools.skill_manager_tool" in lock_source, ( + "tools.skill_manager_tool must still be referenced inside `_ENV_LOCK` " + "for attribute patching (HERMES_HOME / SKILLS_DIR)" + ) + + def test_no_import_statement_for_skill_tools_in_lock(self): + """Double-check: no bare ``import tools.skills_tool`` or + ``import tools.skill_manager_tool`` inside the lock body source.""" + source = _read_streaming() + lines = source.splitlines() + in_lock = False + depth = 0 + for line in lines: + stripped = line.strip() + if stripped.startswith("with _ENV_LOCK:"): + in_lock = True + depth = 0 + continue + if in_lock: + if stripped: + indent = len(line) - len(line.lstrip()) + if depth == 0: + depth = indent + elif indent < depth and stripped: + in_lock = False + continue + # Check for import statements targeting our modules + for mod in _TARGET_MODULES: + # Match both `import tools.skills_tool` and `import tools.skills_tool as _sk` + if f"import {mod}" in stripped: + raise AssertionError( + f"Found `import {mod}` inside `_ENV_LOCK` body — " + f"use sys.modules.get() instead (#2024). Line: {stripped}" + ) \ No newline at end of file