From 587ab683dbee1afb05ab7c7b7862d283d78a2763 Mon Sep 17 00:00:00 2001 From: ethernet Date: Sun, 13 Sep 2026 12:51:17 -0400 Subject: [PATCH] refactor(pm): share worker operation routing and bootstrap policy --- pm/client.py | 26 +++++++------------------- pm/worker.py | 20 +++++--------------- pm/worker_operations.py | 36 ++++++++++++++++++++++++++++++++++++ tests/pm/test_worker.py | 25 ++++++++++++++++++++++--- 4 files changed, 70 insertions(+), 37 deletions(-) create mode 100644 pm/worker_operations.py diff --git a/pm/client.py b/pm/client.py index 16b6998654..976f7269c9 100644 --- a/pm/client.py +++ b/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) diff --git a/pm/worker.py b/pm/worker.py index bd028a0da9..83d42775d5 100644 --- a/pm/worker.py +++ b/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): diff --git a/pm/worker_operations.py b/pm/worker_operations.py new file mode 100644 index 0000000000..2804c257a3 --- /dev/null +++ b/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"), +} diff --git a/tests/pm/test_worker.py b/tests/pm/test_worker.py index 005b35bddf..a602fba83f 100644 --- a/tests/pm/test_worker.py +++ b/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()