refactor(pm): share worker operation routing and bootstrap policy

pm/client.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
 import json
-import os
 from pathlib import Path
 import subprocess
 import threading
@@ -12,6 +11,7 @@ import uuid
 from pm import paths
 from pm.package import InstallError, Runner, StatePackage
 from pm.runtime import is_runtime, runtime_command, runtime_environment
+from pm.worker_operations import OPERATIONS
 
 
 def _members(value):
@@ -41,13 +41,8 @@ def _request(operation, arguments, *, callbacks=None, pause_event=None, project_
     request_id = uuid.uuid4().hex
     update_id = receipt._ambient_update_id()
     callbacks = callbacks or {}
-    names = ([arguments["name"]] if operation in ("ensure", "stage_only") else
-             {"sync_venv": ["venv"], "venv_is_current": ["venv"],
-              "build_environment": ["uv"], "lock_project": ["uv"],
-              "ensure_environment": ["uv"], "ensure_python_tool": ["uv"],
-              "stage_manager_runtime": ["uv"], "check_project_lock": ["uv"],
-              "export_requirements": ["uv"], "build_requirements_environment": ["uv"],
-              "prune_cache": ["uv"]}.get(operation, []))
+    spec = OPERATIONS[operation]
+    names = list(spec.packages) if spec.packages is not None else [arguments["name"]]
     message = {
         "id": request_id, "operation": operation, "arguments": arguments,
         "update_id": update_id,
@@ -60,10 +55,8 @@ def _request(operation, arguments, *, callbacks=None, pause_event=None, project_
     # Bootstrap precedes dispatch and must share the operation's selected cache.
     cache = Path(arguments["cache"]) if arguments.get("cache") is not None else None
     environment = runtime_environment()
-    state_sync = operation in ("sync_venv", "build_environment", "lock_project",
-                               "ensure_environment", "ensure_python_tool", "check_project_lock",
-                               "export_requirements", "build_requirements_environment") or (
-        operation == "ensure" and isinstance(get_package(arguments["name"]), StatePackage))
+    state_sync = spec.bootstrap == "policy" or (
+        spec.bootstrap == "state" and isinstance(get_package(arguments["name"]), StatePackage))
     if (state_sync and not arguments.get("explicit") and not arguments.get("repair")
             and not lazy_installs_allowed()):
         # A ready PM still decides no-op/refusal under its install lock. A cold
@@ -79,7 +72,7 @@ def _request(operation, arguments, *, callbacks=None, pause_event=None, project_
                 receipt.finalize("failed", 1, token=token)
             raise
         environment["HERMES_DISABLE_LAZY_INSTALLS"] = "1"
-    elif operation == "venv_is_current":
+    elif spec.bootstrap == "never":
         command = runtime_command(worker, bootstrap=False, cache=cache)
     else:
         command = runtime_command(worker, cache=cache)
@@ -211,13 +204,8 @@ def stage_only(name, target, *, progress=None) -> Path:
 
 
 def _python_operation(operation: str, arguments: dict):
-    from pm import operations
-    implementation = getattr(operations, operation, None)
-    if implementation is None:
-        from pm import build_operations
-        implementation = getattr(build_operations, operation)
     if is_runtime():
-        return implementation(**arguments)
+        return OPERATIONS[operation].resolve(operation)(**arguments)
     payload = {key: str(value.absolute()) if isinstance(value, Path) else value
                for key, value in arguments.items()}
     return _request(operation, payload)

pm/worker.py

@@ -1,7 +1,6 @@
 """One stdlib JSON-line PM request per isolated process."""
 from __future__ import annotations
 
-import importlib
 import json
 import os
 from pathlib import Path
@@ -74,10 +73,11 @@ def main():
     from pm import paths, receipt
     from pm.package import InstallError
     from pm.registry import load_package_definitions
+    from pm.worker_operations import OPERATIONS
     context = request["context"]
     paths.repo_root = lambda: Path(context["repo"])
     paths.lockfile_path = lambda: Path(context["lockfile"])
-    engine = importlib.import_module("pm.ensure")
+
     call = 0
     callback_lock = threading.Lock()
 
@@ -102,13 +102,8 @@ def main():
     with receipt.worker_context(request.get("update_id")):
         try:
             load_package_definitions(request.get("packages", []))
-            from pm import operations as python
-            operations = {"ensure": engine.ensure, "sync_venv": engine.sync_venv,
-                          "stage_only": engine.stage_only, "venv_is_current": engine.venv_is_current,
-                          "build_environment": python.build_environment, "lock_project": python.lock_project,
-                          "stage_manager_runtime": python.stage_manager_runtime,
-                          "ensure_environment": python.ensure_environment,
-                          "ensure_python_tool": python.ensure_python_tool}
+            operation = request["operation"]
+            implementation = OPERATIONS[operation].resolve(operation)
             arguments = request["arguments"]
             if request["operation"] in ("sync_venv", "venv_is_current"):
                 arguments["plugin_dirs"] = _members(arguments.get("plugin_dirs"))
@@ -117,12 +112,7 @@ def main():
             for name in ("progress", "download_progress"):
                 if name in request["callbacks"]:
                     arguments[name] = lambda *args, name=name: callback(name, *args)
-            if request["operation"] in ("check_project_lock", "export_requirements",
-                                        "build_requirements_environment", "prune_cache"):
-                from pm import build_operations
-                result = getattr(build_operations, request["operation"])(**arguments)
-            else:
-                result = operations[request["operation"]](**arguments)
+            result = implementation(**arguments)
             if request["operation"] == "ensure":
                 result = None  # Runner is reconstructed from the caller's base env.
             if isinstance(result, Path):

pm/worker_operations.py

@@ -0,0 +1,36 @@
+"""Import-light routing and prerequisite policy for the private PM worker.
+
+PM's own runtime is deliberately absent: its public staging entry point must
+run directly, before the worker's dependencies exist.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class Operation:
+    module: str
+    packages: tuple[str, ...] | None  # None selects the request's named package.
+    bootstrap: Literal["always", "policy", "state", "never"]
+
+    def resolve(self, name: str):
+        return getattr(import_module(self.module), name)
+
+
+OPERATIONS = {
+    "ensure": Operation("pm.ensure", None, "state"),
+    "stage_only": Operation("pm.ensure", None, "always"),
+    "sync_venv": Operation("pm.ensure", ("venv",), "policy"),
+    "venv_is_current": Operation("pm.ensure", ("venv",), "never"),
+    "build_environment": Operation("pm.operations", ("uv",), "policy"),
+    "lock_project": Operation("pm.operations", ("uv",), "policy"),
+    "ensure_environment": Operation("pm.operations", ("uv",), "policy"),
+    "ensure_python_tool": Operation("pm.operations", ("uv",), "policy"),
+    "check_project_lock": Operation("pm.build_operations", ("uv",), "policy"),
+    "export_requirements": Operation("pm.build_operations", ("uv",), "policy"),
+    "build_requirements_environment": Operation("pm.build_operations", ("uv",), "policy"),
+    "prune_cache": Operation("pm.build_operations", ("uv",), "always"),
+}

tests/pm/test_worker.py

@@ -602,7 +602,26 @@ def test_worker_side_environment_reuses_and_keeps_selection_on_failed_tool(clien
     assert (root / "active.json").read_bytes() == selection
 
 
-def test_unknown_worker_operation_is_not_dispatched(client):
-    with pytest.raises((KeyError, RuntimeError), match="activate"):
-        client._request("activate", {})
+@pytest.mark.parametrize("operation", ["activate", "stage_manager_runtime"])
+@pytest.mark.parametrize("route", ["client", "wire"])
+def test_unknown_worker_operation_is_not_dispatched(client, monkeypatch, tmp_path, isolated_python, operation, route):
+    import json
+
+    arguments = {"destination": str(tmp_path / "unused")}
+    if route == "client":
+        monkeypatch.setattr(client, "runtime_command", lambda *a, **kw: pytest.fail("unsupported operation acquired PM"))
+        with pytest.raises(KeyError, match=operation):
+            client._request(operation, arguments)
+    else:
+        request = {"id": "unsupported", "operation": operation, "arguments": arguments,
+                   "callbacks": [], "packages": [],
+                   "context": {"repo": str(paths.repo_root()), "lockfile": str(paths.lockfile_path())}}
+        worker = Path(client.__file__).with_name("worker.py")
+        result = subprocess.run(client.runtime_command(worker), input=json.dumps(request) + "\n",
+                                capture_output=True, text=True, encoding="utf-8", timeout=30,
+                                env=client.runtime_environment())
+        assert result.returncode == 0, result.stderr
+        response = json.loads(result.stdout)
+        assert response["error"]["type"] == "KeyError", response
+        assert operation in response["error"]["message"]
     assert not paths.facts_path().exists()