Fix/runner sys path restoration

powermcp/runner.py

@@ -5,22 +5,25 @@
 standalone from a checkout. Three launch styles (declared per tool in the
 registry):
 
-- ``script``: run the entry .py as ``__main__`` (its own ``mcp.run(...)`` fires),
-  after putting the server dir and its parent on ``sys.path`` so existing
-  relative imports (PyPSA's ``sys.path.append`` parent, OpenDSS's
-  ``from core.server import ...``) resolve exactly as in standalone use.
-- ``module``: put the module root on ``sys.path`` and ``runpy.run_module`` the
-  package's ``__main__`` (PSCAD's ``pscad_mcp.main``, HOPE's ``hope_mcp_server``).
-- ``package``: ``runpy.run_module`` a server that ships in its own distribution
-  and is already importable (powerio's ``powerio.mcp``). Nothing goes on
-  ``sys.path`` and this repo bundles no copy of that server.
+- ``script``: run the entry .py as ``__main__`` after temporarily putting
+  the server dir on ``sys.path``.
+- ``module``: put the module root on ``sys.path`` and run its package.
+- ``package``: run a server shipped in its own distribution; no path changes.
+
+All temporary import-path changes are restored when execution returns or
+raises. This matters when PowerMCP is embedded in another Python process;
+standalone CLI execution exits shortly afterwards, but tests and host
+applications may continue importing modules.
 """
 
 from __future__ import annotations
 
 import importlib.util
 import runpy
 import sys
+import threading
+from contextlib import contextmanager
+from collections.abc import Iterator
 
 from .registry import Tool, get_tool, install_hint
 
@@ -29,19 +32,35 @@ class LaunchError(RuntimeError):
     """Raised with an actionable message when a server cannot be launched."""
 
 
-def probe_installed(probe: str | None) -> bool:
-    """True if ``probe`` resolves to a real (importable) module/package.
-
-    Uses find_spec on the FULL dotted name (no module execution). Two cases this
-    must get right:
-    - ``surge`` (top-level): the repo's ``surge/`` folder is a PEP 420 namespace
-      portion (origin None) and must NOT be mistaken for the installed library —
-      so we require a real ``origin``.
-    - ``mhi.pscad`` (subpackage of a namespace): ``mhi`` itself is a namespace
-      package shared by the ``mhi-pscad``/``mhi-psout`` distributions, so probing
-      the top-level ``mhi`` would wrongly look "missing". Probing the full
-      ``mhi.pscad`` resolves the real subpackage (which has an origin).
+_SYS_PATH_LOCK = threading.RLock()
+
+
+@contextmanager
+def _temporary_sys_path(*paths: str) -> Iterator[None]:
+    """Expose paths for one launch and restore the caller's import path.
+
+    Restore the exact caller list contents in place, even if server startup
+    reorders or mutates it. A finally block guarantees restoration when import
+    or server startup raises.
     """
+    # sys.path is process-global. Serialize PowerMCP launches so overlapping
+    # contexts cannot restore snapshots over one another.
+    with _SYS_PATH_LOCK:
+        original = sys.path[:]
+        try:
+            for path in paths:
+                if path not in sys.path:
+                    sys.path.insert(0, path)
+            yield
+        finally:
+            # Servers may reorder, remove, or append entries while importing.
+            # Restore the caller's exact contents in place so references to the
+            # sys.path list remain valid.
+            sys.path[:] = original
+
+
+def probe_installed(probe: str | None) -> bool:
+    """True if ``probe`` resolves to a real (importable) module/package."""
     if not probe:
         return True
     try:
@@ -67,44 +86,30 @@ def _preflight(tool: Tool) -> None:
         raise LaunchError(
             f"{tool.display} requires Windows-only software and cannot run on '{sys.platform}'."
         )
-    # Every server imports the MCP SDK at module scope, so without this the
-    # failure is a raw ImportError traceback rather than an actionable message.
     if not probe_installed("mcp"):
         raise LaunchError(
             "The MCP SDK is not installed; every server needs it.\n"
             "  Install it with:  pip install powermcp"
         )
-    # Only probe pip-provided linchpins. Vendor engines (PSS/E psspy, PSLF) have
-    # probe=None: they live behind a captured path and report their own
-    # actionable error from the server's lazy init.
     if tool.probe and not probe_installed(tool.probe):
         raise LaunchError(
             f"{tool.display}: required package '{tool.probe.split('.')[0]}' is not installed.\n"
-            f"  Install it with:  {install_hint(tool.extra)}"
+            f"  Install it with: {install_hint(tool.extra)}"
         )
 
 
 def _launch_script(tool: Tool) -> None:
     script = tool.resolve_entry_script()
-    # Emulate `python <script>`: only the script's own directory goes on sys.path
-    # (runpy.run_path does not add it). We deliberately do NOT add the parent —
-    # the parent holds dirs named exactly like libraries (pandapower/, surge/),
-    # and each server temporarily exposes the repo root only while importing
-    # shared powermcp modules (OpenDSS also inserts its root for `core`).
     server_dir = str(script.parent)
-    if server_dir not in sys.path:
-        sys.path.insert(0, server_dir)
-    runpy.run_path(str(script), run_name="__main__")
+    with _temporary_sys_path(server_dir):
+        runpy.run_path(str(script), run_name="__main__")
 
 
 def _launch_module(tool: Tool) -> None:
     root = str(tool.resolve_module_root())
-    if root not in sys.path:
-        sys.path.insert(0, root)
-    runpy.run_module(tool.module, run_name="__main__", alter_sys=True)
+    with _temporary_sys_path(root):
+        runpy.run_module(tool.module, run_name="__main__", alter_sys=True)
 
 
 def _launch_package(tool: Tool) -> None:
-    # The server is a dependency, already on sys.path wherever pip put it, so
-    # there is no directory to resolve and nothing of ours to shadow it with.
     runpy.run_module(tool.module, run_name="__main__", alter_sys=True)

tests/test_runner.py

@@ -7,10 +7,13 @@
 from __future__ import annotations
 
 import sys
+import threading
+from types import SimpleNamespace
 
 import pytest
 
 from powermcp import runner
+from powermcp.registry import Tool
 
 
 @pytest.fixture()
@@ -37,16 +40,13 @@ def test_launch_pypsa_runs_once(record_mcp_run):
 
 
 def test_launch_does_not_shadow_installed_library(record_mcp_run):
-    # After launching, `import pandapower` must still resolve to the installed
-    # library (a real package with pandapowerNet), not the server directory.
     runner.launch("pandapower")
     import pandapower as pp
 
     assert hasattr(pp, "create_empty_network")
 
 
 def test_preflight_missing_pip_dep_raises():
-    # andes is an opt-in extra and is not installed in the test venv.
     with pytest.raises(runner.LaunchError) as exc:
         runner.launch("andes")
     assert "powermcp[andes]" in str(exc.value)
@@ -65,26 +65,21 @@ def test_unknown_tool_raises():
 
 
 def test_probe_installed_handles_namespace_parent(tmp_path, monkeypatch):
-    # Reproduces the mhi.pscad case: `ns` is a PEP 420 namespace (no __init__),
-    # `ns.sub` is a real package. Probing the full dotted name must succeed even
-    # though the top-level parent is a namespace; a bare namespace must not.
     (tmp_path / "ns" / "sub").mkdir(parents=True)
     (tmp_path / "ns" / "sub" / "__init__.py").write_text("")
     monkeypatch.syspath_prepend(str(tmp_path))
     for m in ("ns", "ns.sub"):
         sys.modules.pop(m, None)
     try:
-        assert runner.probe_installed("ns.sub") is True   # real subpackage (like mhi.pscad)
-        assert runner.probe_installed("ns") is False       # bare namespace (like the surge shadow)
+        assert runner.probe_installed("ns.sub") is True
+        assert runner.probe_installed("ns") is False
         assert runner.probe_installed("definitely_absent_pkg_xyz") is False
     finally:
         for m in ("ns", "ns.sub"):
             sys.modules.pop(m, None)
 
 
 def test_python_dash_m_powermcp_works():
-    # Regression: MCP client configs launch via `python -m powermcp run <tool>`,
-    # which requires powermcp/__main__.py. Exercise that exact entry path.
     import subprocess
 
     r = subprocess.run(
@@ -96,9 +91,6 @@ def test_python_dash_m_powermcp_works():
 
 
 def test_preflight_missing_sdk_raises(monkeypatch):
-    # Every server imports the MCP SDK at module scope, and a server that dies
-    # during import hands its MCP client nothing at all — the traceback goes to
-    # a stderr the client never reads. So the SDK is checked before launch.
     real = runner.probe_installed
     monkeypatch.setattr(
         runner, "probe_installed", lambda p: False if p == "mcp" else real(p)
@@ -120,3 +112,123 @@ def spy(mod, **kwargs):
     runner.launch("powerio")
     assert seen == ["powerio.mcp"]
     assert len(record_mcp_run) == 1
+
+
+def _tool(kind: str, *, server_dir: str = "example", module_root_rel: str | None = None) -> Tool:
+    return Tool(
+        name=f"test-{kind}", display="test", kind="open-source",
+        windows_only=False, extra=None, server_dir=server_dir,
+        run_kind=kind, entry_rel="server.py" if kind == "script" else None,
+        module="example.main" if kind == "module" else ("powerio.mcp" if kind == "package" else None),
+        module_root_rel=module_root_rel,
+    )
+
+
+@pytest.mark.parametrize("kind", ["script", "module", "package"])
+def test_temporary_import_paths_are_restored_after_success(monkeypatch, kind, tmp_path):
+    original = list(sys.path)
+    temp = str(tmp_path)
+    monkeypatch.setattr(runner, "get_tool", lambda _name: _tool(kind))
+    monkeypatch.setattr(runner, "_preflight", lambda _tool: None)
+    monkeypatch.setattr(runner.Tool, "resolve_entry_script", lambda _self: SimpleNamespace(parent=tmp_path))
+    monkeypatch.setattr(runner.Tool, "resolve_module_root", lambda _self: tmp_path)
+    monkeypatch.setattr(runner.runpy, "run_path", lambda *_a, **_k: None)
+    monkeypatch.setattr(runner.runpy, "run_module", lambda *_a, **_k: None)
+
+    runner.launch("ignored")
+    assert sys.path == original
+    if kind == "package":
+        assert temp not in sys.path
+
+
+@pytest.mark.parametrize("kind", ["script", "module"])
+def test_temporary_import_paths_are_restored_when_launch_raises(monkeypatch, kind, tmp_path):
+    original = list(sys.path)
+    monkeypatch.setattr(runner, "get_tool", lambda _name: _tool(kind))
+    monkeypatch.setattr(runner, "_preflight", lambda _tool: None)
+    monkeypatch.setattr(runner.Tool, "resolve_entry_script", lambda _self: SimpleNamespace(parent=tmp_path))
+    monkeypatch.setattr(runner.Tool, "resolve_module_root", lambda _self: tmp_path)
+
+    def fail(*_args, **_kwargs):
+        raise RuntimeError("simulated startup failure")
+
+    monkeypatch.setattr(runner.runpy, "run_path", fail)
+    monkeypatch.setattr(runner.runpy, "run_module", fail)
+    with pytest.raises(RuntimeError, match="simulated startup failure"):
+        runner.launch("ignored")
+    assert sys.path == original
+
+
+@pytest.mark.parametrize(
+    "mutation",
+    [
+        lambda path: sys.path.append("server-added-entry"),
+        lambda path: sys.path.remove(path),
+        lambda path: sys.path.insert(0, "server-reordered-entry"),
+        lambda path: sys.path.append(path),
+    ],
+)
+def test_temporary_sys_path_restores_exact_state_after_server_mutation(tmp_path, mutation):
+    original = sys.path[:]
+    entry = str(tmp_path)
+    try:
+        with runner._temporary_sys_path(entry):
+            mutation(entry)
+        assert sys.path == original
+    finally:
+        sys.path[:] = original
+
+
+def test_temporary_sys_path_serializes_overlapping_launches(tmp_path):
+    original = sys.path[:]
+    first_entered = threading.Event()
+    release_first = threading.Event()
+    second_started = threading.Event()
+    second_entered = threading.Event()
+    first_path = str(tmp_path / "first")
+    second_path = str(tmp_path / "second")
+
+    def first():
+        with runner._temporary_sys_path(first_path):
+            first_entered.set()
+            assert release_first.wait(timeout=2)
+
+    def second():
+        assert first_entered.wait(timeout=2)
+        second_started.set()
+        with runner._temporary_sys_path(second_path):
+            second_entered.set()
+
+    first_thread = threading.Thread(target=first)
+    second_thread = threading.Thread(target=second)
+    try:
+        first_thread.start()
+        assert first_entered.wait(timeout=2)
+        second_thread.start()
+        assert second_started.wait(timeout=2)
+        assert not second_entered.wait(timeout=0.05)
+        release_first.set()
+        first_thread.join(timeout=2)
+        second_thread.join(timeout=2)
+        assert not first_thread.is_alive()
+        assert not second_thread.is_alive()
+        assert second_entered.is_set()
+        assert sys.path == original
+    finally:
+        release_first.set()
+        first_thread.join(timeout=2)
+        if second_thread.ident is not None:
+            second_thread.join(timeout=2)
+        sys.path[:] = original
+
+
+def test_temporary_sys_path_preserves_preexisting_entries(tmp_path):
+    entry = str(tmp_path)
+    original = list(sys.path)
+    sys.path.insert(0, entry)
+    try:
+        with runner._temporary_sys_path(entry):
+            assert sys.path.count(entry) == 1
+        assert sys.path == [entry, *original]
+    finally:
+        sys.path[:] = original