From bc4688d5837c5252ba398583cccbc755a93ab7ad Mon Sep 17 00:00:00 2001 From: Jack McCarthy Date: Sat, 26 Sep 2026 14:01:01 -0400 Subject: [PATCH 1/5] Switch public REPL execution to disposable workers --- .github/workflows/tests.yml | 20 + servers/README.md | 44 +- servers/lean_client.py | 399 ++++- servers/lean_runtime.py | 677 ++++++-- servers/lsp/server.py | 7 +- servers/repl/core.py | 217 ++- servers/repl/pool.py | 230 ++- servers/repl/server.py | 8 +- tests/fixtures/repl-smoke/.gitignore | 1 + tests/fixtures/repl-smoke/lake-manifest.json | 16 + tests/fixtures/repl-smoke/lakefile.toml | 7 + tests/fixtures/repl-smoke/lean-toolchain | 1 + tests/test_real_repl.py | 38 + tests/test_repl_core_protocol.py | 273 ++- tests/test_repl_pool_lifecycle.py | 390 ++++- tests/test_shared_lean_runtime.py | 1578 +++++++++++++++++- 16 files changed, 3518 insertions(+), 388 deletions(-) create mode 100644 tests/fixtures/repl-smoke/.gitignore create mode 100644 tests/fixtures/repl-smoke/lake-manifest.json create mode 100644 tests/fixtures/repl-smoke/lakefile.toml create mode 100644 tests/fixtures/repl-smoke/lean-toolchain create mode 100644 tests/test_real_repl.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0a2fbdaf..2387eb5d 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -24,3 +24,23 @@ jobs: - run: uv run ruff check autoform_cli servers tests - run: uv run pytest -q - run: make check-example + + real-repl: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 + with: + version: "0.12.1" + python-version: "3.13" + enable-cache: true + - uses: leanprover/lean-action@50fcf42d2e460296f1a34b402e990d1b24f8b596 # v1 + with: + auto-config: false + build: true + build-args: repl + lake-package-directory: tests/fixtures/repl-smoke + - run: uv sync --extra dev --extra repl + - run: uv run pytest -q tests/test_real_repl.py + env: + AUTOFORM_RUN_REAL_REPL_TESTS: "1" diff --git a/servers/README.md b/servers/README.md index e138c07a..eabf24a5 100644 --- a/servers/README.md +++ b/servers/README.md @@ -10,16 +10,24 @@ when useful. Plugin hosts start the two stdio MCP processes automatically. They are lightweight adapters: the first Lean tool call race-safely starts a detached runtime for the current AutoformBot installation, Unix user, and compute node. -That runtime owns one resident REPL pool and LSP session per active Lean -project, so sessions using that installation reuse the same warmed processes. -Closing the session that started it does not stop it; after a crash, the next -tool call starts it again. Runtime sockets include a code fingerprint, so an -in-place upgrade gracefully replaces the older build. +That runtime owns one REPL admission pool and LSP session per active Lean +project. Every public REPL call starts a fresh child and reaps it before +returning, so environments, proof states, and stream contents cannot cross +independent successful calls. If cleanup cannot be confirmed, Autoform returns +an explicit no-replay error, quarantines that project pool, and blocks its +replacement until cleanup succeeds. LSP sessions remain resident because their +protocol is explicitly stateful. Closing the session that started the runtime +does not stop it; after a crash, the next tool call starts it again. Runtime +sockets include a code fingerprint, so an in-place upgrade gracefully replaces +the older build. -REPL and LSP processes remain lazy. A cold tool call stays pending while Lean -warms up, so no `/repl-start`, `/lsp-start`, or model-side sleep is needed. Idle -project processes are closed after 30 minutes by default, while the small -runtime remains available. Its lifecycle is also explicit: +Lean subprocesses remain lazy. A REPL call stays pending while its fresh child +starts, and the first LSP call stays pending while its session starts, so no +`/repl-start`, `/lsp-start`, or model-side sleep is needed. Idle project slots +and LSP sessions are closed after 30 minutes by default, while the small runtime +remains available. The runtime currently requires POSIX process groups and Unix +domain sockets; unsupported platforms fail before starting Lean. Its lifecycle +is also explicit: ```bash uv run autoform-lean-runtime start @@ -30,11 +38,27 @@ uv run autoform-lean-runtime stop `stop` is graceful: it waits for admitted tool calls and Lean children to finish shutting down before a subsequent `start` can replace the runtime. +REPL transport retries are limited to failures detected before the complete +request frame is dispatched. Once the final frame delimiter may have reached +Lean, replay could execute the command twice, so Autoform retires the process +and reports that the outcome is unknown instead of retrying. +The REPL per-call timeout starts before the shared daemon is connected or +started, then covers project admission, fresh child startup, idle-worker wait, +and Lean execution. Verified process cleanup and response delivery get a +separate bounded grace period before the RPC returns. + +`LEAN_REPL_CMD` is a trusted local command. Its descendants must remain in the +dedicated process group Autoform creates; a command that deliberately detaches +with a new session escapes that operating-system cleanup boundary. + The private socket lives below `$XDG_RUNTIME_DIR/autoform`, falling back to a uid-specific directory in `/tmp`; the rotating runtime log is beside it. `AUTOFORM_RUNTIME_DIR` overrides that location. Node-wide limits are controlled by `AUTOFORM_REPL_TOTAL_WORKERS`, `AUTOFORM_REPL_WORKERS_PER_PROJECT`, `AUTOFORM_MAX_LEAN_PROJECTS`, and `AUTOFORM_LEAN_IDLE_SECONDS`. The first process to start the runtime supplies those settings until it is stopped. +`get_repl_status` reports a project pool as `warm` when its admission slots are +cached; it does not mean a Lean REPL child is resident between calls. `AUTOFORM_RUNTIME_RESPONSE_TIMEOUT` can raise the client/daemon response budget -when unusually large worker pools need more than the default 15 minutes to warm. +when a Lean operation and its verified child cleanup need more than the default +15 minutes. diff --git a/servers/lean_client.py b/servers/lean_client.py index 2363a3c8..0a39de06 100644 --- a/servers/lean_client.py +++ b/servers/lean_client.py @@ -7,6 +7,7 @@ from __future__ import annotations +import errno import hashlib import json import math @@ -18,14 +19,30 @@ import time import uuid from dataclasses import dataclass +from logging import getLogger from pathlib import Path from typing import Any +logger = getLogger(__name__) + PROTOCOL_VERSION = 1 MAX_MESSAGE_BYTES = 16 * 1024 * 1024 DEFAULT_CONNECT_TIMEOUT = 2.0 DEFAULT_RESPONSE_TIMEOUT = 900.0 DEFAULT_STARTUP_TIMEOUT = 15.0 +DEFAULT_REPL_REQUEST_TIMEOUT = 180.0 +# The daemon reserves two seconds for verified process cleanup and thirty +# seconds for response/retirement overhead after the public operation deadline. +REPL_RESPONSE_GRACE_SECONDS = 32.0 +_RUNTIME_RESULT_TYPES: dict[str, type[Any]] = { + "daemon.ping": dict, + "daemon.shutdown": dict, + "daemon.status": dict, + "lsp.diagnostics": str, + "lsp.hover": str, + "repl.run": str, + "repl.status": dict, +} PACKAGE_ROOT = Path(__file__).resolve().parent.parent INSTALL_PATH_ID = hashlib.sha256(os.fsencode(PACKAGE_ROOT)).hexdigest()[:10] @@ -52,6 +69,7 @@ def _build_id() -> str: def _build_generation() -> int: """Order in-place builds so an older live wrapper cannot replace a newer one.""" candidates = ( + PACKAGE_ROOT / "servers" / "__init__.py", Path(__file__).resolve(), PACKAGE_ROOT / "servers" / "lean_runtime.py", PACKAGE_ROOT / "servers" / "lsp" / "server.py", @@ -85,10 +103,51 @@ class LeanRuntimeProtocolError(LeanRuntimeError): """The runtime spoke an incompatible or malformed protocol.""" +class LeanRuntimeOutcomeUnknown(LeanRuntimeError): + """A dispatched runtime request did not produce a trustworthy response.""" + + class LeanRuntimeRemoteError(LeanRuntimeError): """The runtime rejected a well-formed request.""" +def _decode_runtime_response(raw: str) -> Any: + """Decode canonical JSON without ambiguous duplicate fields or constants.""" + + def reject_constant(value: str) -> None: + raise LeanRuntimeProtocolError( + f"Lean runtime returned nonstandard JSON constant {value!r}" + ) + + def reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise LeanRuntimeProtocolError( + f"Lean runtime returned duplicate JSON key {key!r}" + ) + result[key] = value + return result + + def parse_finite_float(value: str) -> float: + result = float(value) + if not math.isfinite(result): + raise LeanRuntimeProtocolError( + f"Lean runtime returned non-finite JSON number {value!r}" + ) + return result + + try: + return json.loads( + raw, + parse_constant=reject_constant, + parse_float=parse_finite_float, + object_pairs_hook=reject_duplicate_keys, + ) + except (ValueError, RecursionError) as error: + raise LeanRuntimeProtocolError("Lean runtime returned invalid JSON") from error + + def _response_timeout_from_environment() -> float: raw = os.environ.get( "AUTOFORM_RUNTIME_RESPONSE_TIMEOUT", @@ -107,6 +166,32 @@ def _response_timeout_from_environment() -> float: return value +def _repl_timeout_from_params(params: dict[str, Any]) -> float: + """Resolve the public REPL budget before any daemon startup work begins.""" + timeout = params.get("timeout") + if timeout is None: + raw = os.environ.get( + "AUTOFORM_REPL_REQUEST_TIMEOUT", + str(DEFAULT_REPL_REQUEST_TIMEOUT), + ) + raw = str(DEFAULT_REPL_REQUEST_TIMEOUT) if not raw.strip() else raw + try: + timeout = float(raw) + except ValueError as error: + raise ValueError( + "AUTOFORM_REPL_REQUEST_TIMEOUT must be a number, " + f"got {raw!r}" + ) from error + if ( + isinstance(timeout, bool) + or not isinstance(timeout, (int, float)) + or not math.isfinite(timeout) + or timeout <= 0 + ): + raise ValueError("timeout must be a finite positive number or null") + return float(timeout) + + @dataclass(frozen=True) class RuntimePaths: """Filesystem locations used by one per-user runtime instance.""" @@ -233,32 +318,79 @@ def request( *, autostart: bool | None = None, response_timeout: float | None = None, + deadline: float | None = None, ) -> Any: """Call a runtime method, starting the daemon only before dispatch.""" + request_params = params or {} + if method == "repl.run": + now = time.monotonic() + if deadline is not None and ( + isinstance(deadline, bool) + or not isinstance(deadline, (int, float)) + or not math.isfinite(deadline) + ): + raise LeanRuntimeError( + "Lean runtime request deadline must be a finite number" + ) + if deadline is None: + operation_budget = _repl_timeout_from_params(request_params) + deadline = now + operation_budget + else: + operation_budget = max(0.0, deadline - now) + configured_response_timeout = ( + self.response_timeout + if response_timeout is None + else response_timeout + ) + if ( + isinstance(configured_response_timeout, bool) + or not isinstance(configured_response_timeout, (int, float)) + or not math.isfinite(configured_response_timeout) + or configured_response_timeout <= 0 + ): + raise LeanRuntimeError( + "Lean runtime response timeout must be a finite positive number" + ) + if ( + configured_response_timeout + <= operation_budget + REPL_RESPONSE_GRACE_SECONDS + ): + raise LeanRuntimeError( + "Lean runtime response timeout must exceed the REPL operation " + "timeout plus cleanup grace" + ) should_start = self.autostart if autostart is None else autostart try: return self._request_once( method, - params or {}, + request_params, response_timeout=response_timeout, + deadline=deadline, ) except LeanRuntimeUnavailable: if not should_start: raise - self.ensure_running() + self.ensure_running(deadline=deadline) return self._request_once( method, - params or {}, + request_params, response_timeout=response_timeout, + deadline=deadline, ) - def ping(self, *, autostart: bool = False) -> dict[str, Any]: + def ping( + self, + *, + autostart: bool = False, + deadline: float | None = None, + ) -> dict[str, Any]: """Return daemon identity without warming a Lean project.""" result = self.request( "daemon.ping", autostart=autostart, response_timeout=min(self.response_timeout, 5.0), + deadline=deadline, ) if not isinstance(result, dict): raise LeanRuntimeProtocolError("daemon.ping returned a non-object result") @@ -268,10 +400,26 @@ def ping(self, *, autostart: bool = False) -> dict[str, Any]: ) return result - def ensure_running(self) -> dict[str, Any]: + def ensure_running(self, *, deadline: float | None = None) -> dict[str, Any]: """Race-safely start one detached runtime for this user and node.""" + startup_deadline = time.monotonic() + self.startup_timeout + deadline = ( + startup_deadline + if deadline is None + else min(deadline, startup_deadline) + ) + + def remaining(purpose: str) -> float: + value = deadline - time.monotonic() + if value <= 0: + raise LeanRuntimeUnavailable( + f"timed out waiting for {purpose}; a previous Lean runtime " + "may still be cleaning up" + ) + return value + try: - return self.ping(autostart=False) + return self.ping(autostart=False, deadline=deadline) except LeanRuntimeUnavailable: pass @@ -280,16 +428,31 @@ def ensure_running(self) -> dict[str, Any]: except ImportError as error: # pragma: no cover - guarded by AF_UNIX above raise LeanRuntimeError("runtime bootstrap requires POSIX file locking") from error + def acquire_lock(fd: int, *, purpose: str) -> None: + delay = 0.025 + while True: + wait = remaining(purpose) + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return + except BlockingIOError: + pass + except OSError as error: + if error.errno != errno.EACCES: + raise + time.sleep(min(delay, wait)) + delay = min(delay * 1.7, 0.25) + _private_runtime_directory(self.paths.directory) lock_fd = os.open(self.paths.lock, os.O_CREAT | os.O_RDWR, 0o600) try: - fcntl.flock(lock_fd, fcntl.LOCK_EX) + acquire_lock(lock_fd, purpose="Lean runtime startup coordination") try: - return self.ping(autostart=False) + return self.ping(autostart=False, deadline=deadline) except LeanRuntimeUnavailable: pass - self._stop_previous_builds() + self._stop_previous_builds(deadline=deadline) # A daemon owns this lock for its complete lifetime. If it has # stopped accepting connections but is still draining requests, @@ -300,24 +463,27 @@ def ensure_running(self) -> dict[str, Any]: 0o600, ) try: - fcntl.flock(lifetime_fd, fcntl.LOCK_EX) + acquire_lock(lifetime_fd, purpose="the previous Lean runtime") + remaining("Lean runtime startup") self._remove_stale_socket() process = self._spawn_daemon() finally: os.close(lifetime_fd) try: - deadline = time.monotonic() + self.startup_timeout delay = 0.025 last_error: BaseException | None = None while time.monotonic() < deadline: if process.poll() is not None: break try: - return self.ping(autostart=False) + return self.ping(autostart=False, deadline=deadline) except LeanRuntimeUnavailable as error: last_error = error - time.sleep(delay) + wait = deadline - time.monotonic() + if wait <= 0: + break + time.sleep(min(delay, wait)) delay = min(delay * 1.7, 0.25) exit_detail = ( @@ -336,31 +502,36 @@ def ensure_running(self) -> dict[str, Any]: finally: os.close(lock_fd) - def stop(self) -> dict[str, Any]: + def stop(self, *, deadline: float | None = None) -> dict[str, Any]: """Ask a running daemon to finish active calls and shut down.""" try: result = self.request( "daemon.shutdown", autostart=False, response_timeout=10.0, + deadline=deadline, ) except LeanRuntimeUnavailable: - stopped = self._stop_previous_builds() + stopped = self._stop_previous_builds(deadline=deadline) if stopped: return {"stopping": False, "stopped_previous": stopped} raise if not isinstance(result, dict): raise LeanRuntimeProtocolError("daemon.shutdown returned a non-object result") - deadline = time.monotonic() + self.response_timeout - while self.paths.socket.exists() and time.monotonic() < deadline: - time.sleep(0.025) + stop_deadline = ( + time.monotonic() + self.response_timeout + if deadline is None + else deadline + ) + while self.paths.socket.exists() and time.monotonic() < stop_deadline: + time.sleep(min(0.025, max(0.0, stop_deadline - time.monotonic()))) if self.paths.socket.exists(): raise LeanRuntimeError( f"Lean runtime is still draining requests at {self.paths.socket}" ) return result - def _stop_previous_builds(self) -> list[int]: + def _stop_previous_builds(self, *, deadline: float | None = None) -> list[int]: """Gracefully replace older code generations at the same install path.""" if not self._uses_default_paths: return [] @@ -377,7 +548,11 @@ def _stop_previous_builds(self) -> list[int]: startup_timeout=self.startup_timeout, ) try: - status = previous.request("daemon.ping", autostart=False) + status = previous.request( + "daemon.ping", + autostart=False, + deadline=deadline, + ) except LeanRuntimeUnavailable: continue generation = ( @@ -388,7 +563,7 @@ def _stop_previous_builds(self) -> list[int]: "a newer Autoform runtime build is already active; " "restart this plugin session before using Lean tools" ) - result = previous.stop() + result = previous.stop(deadline=deadline) pid = result.get("pid") if isinstance(result, dict) else None if isinstance(pid, int): stopped.append(pid) @@ -403,6 +578,8 @@ def _spawn_daemon(self) -> subprocess.Popen[bytes]: str(self.paths.socket), "--log", str(self.paths.log), + "--lifetime-lock", + str(self.paths.lifetime_lock), "serve", ] with self.paths.log.open("ab", buffering=0) as log: @@ -433,14 +610,22 @@ def _remove_stale_socket(self) -> None: @staticmethod def _terminate_failed_start(process: subprocess.Popen[bytes]) -> None: + """Request shutdown without killing a daemon that may own active work.""" if process.poll() is not None: return - process.terminate() try: + process.terminate() process.wait(timeout=5) except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=5) + # The daemon may already be serving another client and draining a + # Lean child. Its lifetime lock prevents a replacement from being + # admitted while fail-closed cleanup continues. + logger.warning( + "spawned Lean runtime is still shutting down; leaving it under " + "its lifetime lock" + ) + except OSError: + logger.exception("failed to request shutdown of the spawned Lean runtime") def _request_once( self, @@ -448,24 +633,43 @@ def _request_once( params: dict[str, Any], *, response_timeout: float | None = None, + deadline: float | None = None, ) -> Any: request_id = uuid.uuid4().hex - payload = json.dumps( - { - "v": PROTOCOL_VERSION, - "id": request_id, - "method": method, - "params": params, - }, - separators=(",", ":"), - ).encode("utf-8") + b"\n" + request: dict[str, Any] = { + "v": PROTOCOL_VERSION, + "id": request_id, + "method": method, + "params": params, + } + if method == "repl.run" and deadline is not None: + request["deadline"] = deadline + payload = json.dumps(request, separators=(",", ":")).encode("utf-8") + b"\n" if len(payload) > MAX_MESSAGE_BYTES: raise LeanRuntimeProtocolError("Lean runtime request exceeds the message limit") connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) dispatched = False + + def outcome_unknown() -> LeanRuntimeOutcomeUnknown: + return LeanRuntimeOutcomeUnknown( + "Lean runtime request may have completed, but no trustworthy " + "response was received after request dispatch. The request was " + "not retried and must not be replayed." + ) + + def bounded_timeout(configured: float, *, grace: float = 0.0) -> float: + if deadline is None: + return configured + remaining = deadline + grace - time.monotonic() + if remaining <= 0: + raise LeanRuntimeUnavailable( + "Lean runtime request deadline expired before dispatch" + ) + return min(configured, remaining) + try: - connection.settimeout(self.connect_timeout) + connection.settimeout(bounded_timeout(self.connect_timeout)) try: connection.connect(str(self.paths.socket)) except (FileNotFoundError, ConnectionRefusedError) as error: @@ -479,60 +683,133 @@ def _request_once( ) from error raise LeanRuntimeError(f"cannot connect to Lean runtime: {error}") from error - connection.settimeout(response_timeout or self.response_timeout) + configured_response_timeout = ( + self.response_timeout + if response_timeout is None + else response_timeout + ) + connection.settimeout(bounded_timeout(configured_response_timeout)) + if deadline is not None and time.monotonic() >= deadline: + raise LeanRuntimeUnavailable( + "Lean runtime request deadline expired before dispatch" + ) # From this point onward, any failure is ambiguous: the daemon may # have received the request. Never auto-replay Lean execution. dispatched = True connection.sendall(payload) - raw = self._read_line(connection) + response_grace = ( + REPL_RESPONSE_GRACE_SECONDS + if method == "repl.run" and deadline is not None + else 0.0 + ) + configured_response_deadline = ( + time.monotonic() + configured_response_timeout + ) + response_deadline = configured_response_deadline + if deadline is not None: + response_deadline = min( + response_deadline, + deadline + response_grace, + ) + raw = self._read_line(connection, deadline=response_deadline) except socket.timeout as error: - phase = "response" if dispatched else "connection" - raise LeanRuntimeError(f"timed out waiting for Lean runtime {phase}") from error - except LeanRuntimeError: + if dispatched: + raise outcome_unknown() from error + raise LeanRuntimeError("timed out waiting for Lean runtime connection") from error + except LeanRuntimeError as error: + if dispatched: + raise outcome_unknown() from error raise except OSError as error: if not dispatched: raise LeanRuntimeUnavailable( f"Lean runtime is not listening at {self.paths.socket}" ) from error - raise LeanRuntimeError( - "connection to Lean runtime closed after request dispatch; the request was not retried" - ) from error + raise outcome_unknown() from error finally: connection.close() try: - response = json.loads(raw) - except (UnicodeDecodeError, json.JSONDecodeError) as error: - raise LeanRuntimeProtocolError("Lean runtime returned invalid JSON") from error - if not isinstance(response, dict): - raise LeanRuntimeProtocolError("Lean runtime response is not an object") - if response.get("v") != PROTOCOL_VERSION: - raise LeanRuntimeProtocolError( - f"Lean runtime protocol mismatch: expected {PROTOCOL_VERSION}, got {response.get('v')!r}" + response = _decode_runtime_response(raw) + if not isinstance(response, dict): + raise LeanRuntimeProtocolError("Lean runtime response is not an object") + version = response.get("v") + if type(version) is not int or version != PROTOCOL_VERSION: + raise LeanRuntimeProtocolError( + "Lean runtime protocol mismatch: expected " + f"{PROTOCOL_VERSION}, got {version!r}" + ) + response_id = response.get("id") + if type(response_id) is not str or response_id != request_id: + raise LeanRuntimeProtocolError( + "Lean runtime response id does not match the request" + ) + ok = response.get("ok") + if type(ok) is not bool: + raise LeanRuntimeProtocolError( + "Lean runtime response has an invalid success marker" + ) + expected_keys = ( + {"v", "id", "ok", "result"} + if ok + else {"v", "id", "ok", "error"} ) - if response.get("id") != request_id: - raise LeanRuntimeProtocolError("Lean runtime response id does not match the request") - if response.get("ok") is True: - return response.get("result") - error = response.get("error") - if not isinstance(error, dict): - raise LeanRuntimeProtocolError("Lean runtime returned a malformed error") - error_type = error.get("type", "RuntimeError") - message = error.get("message", "unspecified runtime error") + if set(response) != expected_keys: + raise LeanRuntimeProtocolError( + "Lean runtime response has an invalid envelope" + ) + if ok: + expected_result_type = _RUNTIME_RESULT_TYPES.get(method) + if expected_result_type is None: + raise LeanRuntimeProtocolError( + f"Lean runtime returned success for unknown method {method!r}" + ) + result = response["result"] + if not isinstance(result, expected_result_type): + raise LeanRuntimeProtocolError( + f"Lean runtime {method} returned an invalid result type" + ) + return result + remote_error = response["error"] + if not isinstance(remote_error, dict) or set(remote_error) != { + "type", + "message", + }: + raise LeanRuntimeProtocolError("Lean runtime returned a malformed error") + error_type = remote_error["type"] + message = remote_error["message"] + if type(error_type) is not str or type(message) is not str: + raise LeanRuntimeProtocolError("Lean runtime returned a malformed error") + except LeanRuntimeProtocolError as error: + raise outcome_unknown() from error raise LeanRuntimeRemoteError(f"{error_type}: {message}") @staticmethod - def _read_line(connection: socket.socket) -> str: + def _read_line(connection: socket.socket, *, deadline: float) -> str: data = bytearray() while len(data) <= MAX_MESSAGE_BYTES: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise socket.timeout("Lean runtime response deadline expired") + connection.settimeout(remaining) chunk = connection.recv(min(65536, MAX_MESSAGE_BYTES + 1 - len(data))) + if time.monotonic() >= deadline: + raise socket.timeout("Lean runtime response deadline expired") if not chunk: raise LeanRuntimeProtocolError("Lean runtime closed without a response") data.extend(chunk) + if len(data) > MAX_MESSAGE_BYTES: + raise LeanRuntimeProtocolError( + "Lean runtime response exceeds the message limit" + ) newline = data.find(b"\n") if newline >= 0: if data[newline + 1 :]: raise LeanRuntimeProtocolError("Lean runtime returned trailing response data") - return bytes(data[:newline]).decode("utf-8") + try: + return bytes(data[:newline]).decode("utf-8") + except UnicodeDecodeError as error: + raise LeanRuntimeProtocolError( + "Lean runtime returned invalid UTF-8" + ) from error raise LeanRuntimeProtocolError("Lean runtime response exceeds the message limit") diff --git a/servers/lean_runtime.py b/servers/lean_runtime.py index b328c797..6d72dc38 100644 --- a/servers/lean_runtime.py +++ b/servers/lean_runtime.py @@ -28,10 +28,12 @@ from servers import resolve_lean_file, resolve_lean_project_dir from servers.lean_client import ( BUILD_GENERATION, + DEFAULT_REPL_REQUEST_TIMEOUT, INSTALL_ID, MAX_MESSAGE_BYTES, PROTOCOL_VERSION, DEFAULT_RESPONSE_TIMEOUT, + REPL_RESPONSE_GRACE_SECONDS, LeanRuntimeClient, LeanRuntimeError, LeanRuntimeUnavailable, @@ -46,12 +48,13 @@ LspProtocolError, format_lsp_diagnostics, ) -from servers.repl.core import DEFAULT_REPL_STARTUP_TIMEOUT, format_repl_response +from servers.repl.core import format_repl_response from servers.repl.pool import ( DEFAULT_RAM_FRACTION, - DEFAULT_STARTUP_STAGGER_SECONDS, LeanReplPool, LeanReplPoolConfig, + ReplPoolBusyError, + ReplPoolUnavailableError, ) logger = logging.getLogger(__name__) @@ -62,15 +65,16 @@ DEFAULT_IDLE_SECONDS = 30 * 60 DEFAULT_LSP_TIMEOUT = 60.0 DEFAULT_MAX_LSP_REQUEST_SECONDS = 600.0 -DEFAULT_REPL_REQUEST_TIMEOUT = 30.0 DEFAULT_MAX_REPL_REQUEST_SECONDS = 240.0 DEFAULT_RPC_READ_TIMEOUT = 10.0 DEFAULT_MAX_CONNECTIONS = 64 RUNTIME_SAFETY_SECONDS = 30.0 -# Conservative bounds for cleanup/startup work that surrounds one tool call. +TERMINAL_CLEANUP_RETRY_SECONDS = 0.05 +MAX_TERMINAL_CLEANUP_RETRY_SECONDS = 1.0 +MAX_REQUEST_ID_CHARS = 128 +# Conservative bounds for cleanup/startup work that surrounds one LSP call. # They keep the daemon's work inside the client's response deadline even when # an inactive project must be replaced first. -REPL_WORKER_CLOSE_BUDGET = 10.0 LSP_STARTUP_BUDGET = 60.0 LSP_CLOSE_BUDGET = 65.0 @@ -79,13 +83,35 @@ class ProjectResourceBusyError(TimeoutError): """A shared project slot could not be admitted within the RPC budget.""" -def _repl_creation_budget(worker_count: int) -> float: - """Bound victim cleanup, cold startup, and failed-start cleanup.""" - return ( - worker_count * DEFAULT_REPL_STARTUP_TIMEOUT - + max(0, worker_count - 1) * DEFAULT_STARTUP_STAGGER_SECONDS - + 2 * worker_count * REPL_WORKER_CLOSE_BUDGET - ) +def _decode_runtime_request(raw: bytes) -> Any: + """Decode strict JSON for the runtime's private request boundary.""" + + def reject_constant(value: str) -> None: + raise ValueError(f"request contains nonstandard JSON constant {value!r}") + + def reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"request contains duplicate JSON key {key!r}") + result[key] = value + return result + + def parse_finite_float(value: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError(f"request contains non-finite JSON number {value!r}") + return result + + try: + return json.loads( + raw, + parse_constant=reject_constant, + parse_float=parse_finite_float, + object_pairs_hook=reject_duplicate_keys, + ) + except (UnicodeDecodeError, ValueError, RecursionError) as error: + raise ValueError("request is not valid strict JSON") from error def _positive_int(name: str, default: int) -> int: @@ -209,16 +235,14 @@ def from_environment(cls) -> "LeanRuntimeConfig": "AUTOFORM_RUNTIME_RESPONSE_TIMEOUT", DEFAULT_RESPONSE_TIMEOUT, ) - repl_creation_budget = _repl_creation_budget(workers_per_project) if ( - repl_creation_budget - + max_repl_request_seconds - + RUNTIME_SAFETY_SECONDS - > response_timeout + max_repl_request_seconds + + REPL_RESPONSE_GRACE_SECONDS + >= response_timeout ): raise ValueError( "AUTOFORM_RUNTIME_RESPONSE_TIMEOUT is too small for the configured " - "REPL worker startup and request limits" + "REPL request and cleanup limits" ) if ( LSP_CLOSE_BUDGET @@ -319,6 +343,8 @@ def __init__( self._is_valid = is_valid self._clock = clock self._entries: dict[Path, _CacheEntry[T]] = {} + self._retiring: dict[Path, T] = {} + self._retiring_active: set[Path] = set() self._creating: set[Path] = set() self._condition = threading.Condition() self._closed = False @@ -341,21 +367,44 @@ def lease( *, create: bool = True, acquisition_timeout: float | None = None, + deadline: float | None = None, creation_budget: float = 0.0, ) -> Iterator[T | None]: """Keep a project resource alive for the complete operation.""" + if acquisition_timeout is not None and deadline is not None: + raise TypeError("pass acquisition_timeout or deadline, not both") + if acquisition_timeout is not None: + if acquisition_timeout <= 0: + raise ProjectResourceBusyError( + "no response budget remains for a shared Lean project slot" + ) + deadline = self._clock() + acquisition_timeout root = resolve_lean_project_dir(project_dir) resource = self._acquire( root, create=create, - acquisition_timeout=acquisition_timeout, + deadline=deadline, creation_budget=creation_budget, ) + operation_error: BaseException | None = None try: yield resource - finally: + except BaseException as error: + operation_error = error + try: if resource is not None: self._release(root, resource) + except BaseException as cleanup_error: + if operation_error is None: + raise + note = f"Lean project resource release also failed: {cleanup_error}" + add_note = getattr(operation_error, "add_note", None) + if add_note is not None: + add_note(note) + else: # pragma: no cover - Python 3.10 compatibility + logger.error("%s", note) + if operation_error is not None: + raise operation_error.with_traceback(operation_error.__traceback__) def stats(self) -> dict[str, Any]: with self._condition: @@ -366,22 +415,35 @@ def stats(self) -> dict[str, Any]: { "project_dir": str(root), "active": entry.active, - "valid": not entry.invalid, + "valid": ( + not entry.invalid + and ( + self._is_valid is None + or self._is_valid(entry.resource) + ) + ), "idle_seconds": round(max(0.0, now - entry.last_used), 3), } for root, entry in sorted( self._entries.items(), key=lambda item: str(item[0]) ) ], + "retiring": sorted(str(root) for root in self._retiring), "creating": sorted(str(root) for root in self._creating), } def state(self, project_dir: str) -> str: - """Return ``cold``, ``warming``, or ``warm`` without creating state.""" + """Return the current project-resource lifecycle state.""" root = resolve_lean_project_dir(project_dir) with self._condition: - if root in self._entries: - return "warm" + entry = self._entries.get(root) + if entry is not None: + valid = not entry.invalid and ( + self._is_valid is None or self._is_valid(entry.resource) + ) + return "warm" if valid else "retiring" + if root in self._retiring: + return "retiring" if root in self._creating: return "warming" return "cold" @@ -400,84 +462,103 @@ def evict_idle(self) -> int: if self._idle_seconds <= 0: return 0 with self._condition: + if self._closed: + return 0 now = self._clock() victims = [ root for root, entry in self._entries.items() if entry.active == 0 and now - entry.last_used >= self._idle_seconds ] - resources = [self._entries.pop(root).resource for root in victims] + for root in victims: + self._retiring[root] = self._entries.pop(root).resource + retiring = [ + (root, resource) + for root, resource in self._retiring.items() + if root not in self._retiring_active + ] if victims: self._condition.notify_all() - self._close_many(resources) - return len(resources) + for root, resource in retiring: + with self._condition: + if ( + self._retiring.get(root) is not resource + or root in self._retiring_active + ): + continue + self._retiring_active.add(root) + self._retire(root, resource) + return len(victims) def close(self) -> None: """Stop admission, wait for active leases, then close all resources.""" self._stop_sweeper.set() with self._condition: self._closed = True - while self._creating or any(entry.active for entry in self._entries.values()): + while ( + self._creating + or self._retiring_active + or any(entry.active for entry in self._entries.values()) + ): self._condition.wait(timeout=0.5) - resources = [entry.resource for entry in self._entries.values()] + for root, entry in self._entries.items(): + self._retiring[root] = entry.resource self._entries.clear() + retiring = list(self._retiring.items()) self._condition.notify_all() - self._close_many(resources) - if self._sweeper and self._sweeper is not threading.current_thread(): - self._sweeper.join() + first_error: BaseException | None = None + try: + for root, resource in retiring: + with self._condition: + self._condition.wait_for( + lambda: root not in self._retiring_active + ) + if self._retiring.get(root) is not resource: + continue + self._retiring_active.add(root) + try: + self._retire(root, resource) + except BaseException as error: + if first_error is None: + first_error = error + finally: + if self._sweeper and self._sweeper is not threading.current_thread(): + self._sweeper.join() + if first_error is not None: + raise first_error.with_traceback(first_error.__traceback__) + with self._condition: + failed = len(self._retiring) + if failed: + raise RuntimeError( + f"failed to retire {failed} Lean project resource(s)" + ) def _acquire( self, root: Path, *, create: bool, - acquisition_timeout: float | None, + deadline: float | None, creation_budget: float, ) -> T | None: - if acquisition_timeout is not None and acquisition_timeout <= 0: + if deadline is not None and self._clock() >= deadline: raise ProjectResourceBusyError( "no response budget remains for a shared Lean project slot" ) if creation_budget < 0: raise ValueError("creation_budget must be nonnegative") fingerprint = lean_project_fingerprint(root) - deadline = ( - self._clock() + acquisition_timeout - if acquisition_timeout is not None - else None - ) - resources_to_close: list[T] = [] - reserved = False - while True: wait = False + retirement: tuple[Path, T] | None = None with self._condition: if self._closed: raise RuntimeError("project resource cache is closed") - entry = self._entries.get(root) - entry_is_stale = ( - entry is not None - and ( - entry.invalid - or entry.fingerprint != fingerprint - or ( - self._is_valid is not None - and not self._is_valid(entry.resource) - ) - ) - ) - if entry_is_stale: - assert entry is not None - if entry.active: - if not create: - resource = None - break - self._require_creation_budget( - root, - deadline=deadline, - creation_budget=creation_budget, - ) + if root in self._retiring: + if not create: + return None + if root in self._retiring_active: wait = True else: self._require_creation_budget( @@ -485,34 +566,72 @@ def _acquire( deadline=deadline, creation_budget=creation_budget, ) - resources_to_close.append(self._entries.pop(root).resource) - self._condition.notify_all() - entry = None - - if not wait and entry is not None: + self._retiring_active.add(root) + retirement = (root, self._retiring[root]) + entry = None + else: + entry = self._entries.get(root) + entry_is_stale = ( + entry is not None + and ( + entry.invalid + or entry.fingerprint != fingerprint + or ( + self._is_valid is not None + and not self._is_valid(entry.resource) + ) + ) + ) + if entry_is_stale: + assert entry is not None + if entry.active: + if not create: + return None + self._require_creation_budget( + root, + deadline=deadline, + creation_budget=creation_budget, + ) + wait = True + else: + self._require_creation_budget( + root, + deadline=deadline, + creation_budget=creation_budget, + ) + resource = self._entries.pop(root).resource + self._retiring[root] = resource + self._retiring_active.add(root) + retirement = (root, resource) + self._condition.notify_all() + entry = None + + if retirement is None and not wait and entry is not None: if deadline is not None and self._clock() >= deadline: raise ProjectResourceBusyError( f"timed out waiting for a shared Lean project slot: {root}" ) entry.active += 1 entry.last_used = self._clock() - resource = entry.resource - break + return entry.resource - if not wait and entry is None and not create: - resource = None - break + if retirement is None and not wait and entry is None and not create: + return None - if not wait and root in self._creating: + if retirement is None and not wait and root in self._creating: wait = True - if not wait: + if retirement is None and not wait: self._require_creation_budget( root, deadline=deadline, creation_budget=creation_budget, ) - occupied = len(self._entries) + len(self._creating) + occupied = ( + len(self._entries) + + len(self._creating) + + len(self._retiring) + ) if occupied >= self._max_entries: inactive = [ (candidate.last_used, path) @@ -521,36 +640,59 @@ def _acquire( ] if inactive: _, victim = min(inactive) - resources_to_close.append(self._entries.pop(victim).resource) + resource = self._entries.pop(victim).resource + self._retiring[victim] = resource + self._retiring_active.add(victim) + retirement = (victim, resource) + self._condition.notify_all() + elif any( + path not in self._retiring_active + for path in self._retiring + ): + retiring_root = next( + path + for path in self._retiring + if path not in self._retiring_active + ) + self._retiring_active.add(retiring_root) + retirement = ( + retiring_root, + self._retiring[retiring_root], + ) else: wait = True - if not wait: + if retirement is None and not wait: self._creating.add(root) - reserved = True self._condition.notify_all() - resource = None break - wait_seconds = 0.5 - if deadline is not None: - remaining = deadline - self._clock() - if remaining <= 0: - raise ProjectResourceBusyError( - f"timed out waiting for a shared Lean project slot: {root}" - ) - wait_seconds = min(wait_seconds, remaining) - self._condition.wait(timeout=wait_seconds) - - if resources_to_close: - self._close_many(resources_to_close) - resources_to_close.clear() - - if resources_to_close: - self._close_many(resources_to_close) + if retirement is None: + wait_seconds = 0.5 + if deadline is not None: + remaining = deadline - self._clock() + if remaining <= 0: + raise ProjectResourceBusyError( + f"timed out waiting for a shared Lean project slot: {root}" + ) + wait_seconds = min(wait_seconds, remaining) + self._condition.wait(timeout=wait_seconds) + + if retirement is not None: + retiring_root, retiring_resource = retirement + if not self._retire(retiring_root, retiring_resource): + raise ProjectResourceBusyError( + "failed to retire a stale Lean project resource: " + f"{retiring_root}" + ) - if not reserved: - return resource + if deadline is not None and self._clock() >= deadline: + with self._condition: + self._creating.discard(root) + self._condition.notify_all() + raise ProjectResourceBusyError( + f"shared Lean project startup deadline expired: {root}" + ) try: created = self._factory(root) @@ -576,9 +718,16 @@ def _acquire( last_used=self._clock(), active=1, ) + if close_created: + self._retiring[root] = created + self._retiring_active.add(root) self._condition.notify_all() if close_created: - self._safe_close(created) + cleaned = self._retire(root, created) + if not cleaned: + raise ProjectResourceBusyError( + f"failed to retire a late Lean project resource: {root}" + ) if startup_expired: raise ProjectResourceBusyError( f"shared Lean project startup exceeded its response budget: {root}" @@ -595,19 +744,47 @@ def _require_creation_budget( ) -> None: if deadline is None: return - if deadline - self._clock() < creation_budget: + if deadline - self._clock() <= creation_budget: raise ProjectResourceBusyError( f"not enough response budget to start a shared Lean project slot: {root}" ) def _release(self, root: Path, resource: T) -> None: + retirement: tuple[Path, T] | None = None + validation_error: BaseException | None = None with self._condition: entry = self._entries.get(root) if entry is None or entry.resource is not resource: raise RuntimeError("project resource lease is no longer registered") entry.active -= 1 entry.last_used = self._clock() + invalid = entry.invalid + if not invalid and self._is_valid is not None: + try: + invalid = not self._is_valid(resource) + except BaseException as error: + validation_error = error + invalid = True + if entry.active == 0 and invalid: + self._entries.pop(root) + self._retiring[root] = resource + self._retiring_active.add(root) + retirement = (root, resource) self._condition.notify_all() + if retirement is not None: + try: + self._retire(*retirement) + except BaseException as cleanup_error: + if validation_error is None: + raise + note = f"Lean project resource cleanup also failed: {cleanup_error}" + add_note = getattr(validation_error, "add_note", None) + if add_note is not None: + add_note(note) + else: # pragma: no cover - Python 3.10 compatibility + logger.error("%s", note) + if validation_error is not None: + raise validation_error.with_traceback(validation_error.__traceback__) def _sweep(self, interval: float) -> None: while not self._stop_sweeper.wait(interval): @@ -616,15 +793,23 @@ def _sweep(self, interval: float) -> None: except Exception: logger.exception("failed to evict idle Lean project resources") - def _close_many(self, resources: list[T]) -> None: - for resource in resources: - self._safe_close(resource) - - def _safe_close(self, resource: T) -> None: + def _retire(self, root: Path, resource: T) -> bool: + """Try one bounded close while retaining failed ownership in quarantine.""" + succeeded = False try: self._close_resource(resource) except Exception: logger.exception("failed to close Lean project resource") + return False + else: + succeeded = True + return True + finally: + with self._condition: + self._retiring_active.discard(root) + if succeeded and self._retiring.get(root) is resource: + self._retiring.pop(root) + self._condition.notify_all() class LeanRuntimeServices: @@ -640,9 +825,6 @@ def __init__( ) -> None: self.config = config or LeanRuntimeConfig.from_environment() self.started_at = time.monotonic() - self.repl_creation_budget = _repl_creation_budget( - self.config.repl_workers_per_project - ) self.lsp_creation_budget = LSP_STARTUP_BUDGET + LSP_CLOSE_BUDGET def default_repl_factory(project_dir: Path) -> LeanReplPool: @@ -666,11 +848,15 @@ def default_lsp_factory(project_dir: Path) -> LeanLspSession: session.start() return session + def close_repl_pool(pool: LeanReplPool) -> None: + pool.shutdown() + self.repl_projects = ProjectResourceCache( repl_factory or default_repl_factory, - lambda pool: pool.shutdown(), + close_repl_pool, max_entries=self.config.repl_project_limit, idle_seconds=self.config.idle_seconds, + is_valid=lambda pool: getattr(pool, "is_usable", lambda: True)(), start_sweeper=start_sweepers, ) self.lsp_projects = ProjectResourceCache( @@ -682,7 +868,13 @@ def default_lsp_factory(project_dir: Path) -> LeanLspSession: start_sweeper=start_sweepers, ) - def dispatch(self, method: str, params: dict[str, Any]) -> Any: + def dispatch( + self, + method: str, + params: dict[str, Any], + *, + client_deadline: float | None = None, + ) -> Any: if method == "daemon.ping": return self.status(include_projects=False) if method == "daemon.status": @@ -707,17 +899,79 @@ def dispatch(self, method: str, params: dict[str, Any]) -> Any: "timeout exceeds the node-wide limit of " f"{self.config.max_repl_request_seconds:g} seconds" ) - with self.repl_projects.lease( - project_dir, - acquisition_timeout=self._acquisition_timeout(effective_timeout), - creation_budget=self.repl_creation_budget, - ) as pool: - assert pool is not None - return format_repl_response(pool.run(code, timeout=effective_timeout)) + if client_deadline is not None and ( + isinstance(client_deadline, bool) + or not isinstance(client_deadline, (int, float)) + or not math.isfinite(client_deadline) + ): + raise ValueError("deadline must be a finite number or null") + server_deadline = time.monotonic() + effective_timeout + deadline = ( + server_deadline + if client_deadline is None + else min(server_deadline, float(client_deadline)) + ) + if time.monotonic() >= deadline: + raise ProjectResourceBusyError( + "Lean REPL request deadline expired before admission" + ) + repl_result: Any = None + repl_attempted = False + repl_completed = False + try: + with self.repl_projects.lease( + project_dir, + deadline=deadline, + creation_budget=0.0, + ) as pool: + assert pool is not None + remaining = deadline - time.monotonic() + if remaining <= 0: + raise ProjectResourceBusyError( + "Lean REPL request deadline expired before execution" + ) + repl_attempted = True + repl_result = pool.run(code, deadline=deadline) + repl_completed = True + except (ReplPoolBusyError, ReplPoolUnavailableError): + raise + except Exception as error: + if not repl_attempted: + raise + phase = ( + "project resource release" + if repl_completed + else "REPL execution" + ) + return format_repl_response( + { + "repl_error": ( + f"Lean {phase} failed after the command may have been " + f"dispatched: {error}. The request must not be replayed." + ), + "outcome_unknown": True, + } + ) + try: + return format_repl_response(repl_result) + except Exception as error: + return format_repl_response( + { + "repl_error": ( + "Lean REPL command produced an invalid result: " + f"{error}. The request must not be replayed." + ), + "outcome_unknown": True, + } + ) if method == "repl.status": project_dir = self._string_param(params, "project_dir") with self.repl_projects.lease(project_dir, create=False) as pool: - state = "warm" if pool is not None else self.repl_projects.state(project_dir) + state = ( + "warm" + if pool is not None + else self.repl_projects.state(project_dir) + ) return { "state": state, "capacity": ( @@ -728,7 +982,9 @@ def dispatch(self, method: str, params: dict[str, Any]) -> Any: "memory_usage_gb": ( round(pool.get_memory_usage(), 2) if pool is not None else 0.0 ), - "shutdown": pool._shutdown if pool is not None else False, + "shutdown": ( + pool._shutdown if pool is not None else state == "retiring" + ), "daemon_pid": os.getpid(), "node_total_workers": self.config.total_repl_workers, } @@ -792,8 +1048,23 @@ def status(self, *, include_projects: bool) -> dict[str, Any]: return result def close(self) -> None: - self.repl_projects.close() - self.lsp_projects.close() + repl_error: BaseException | None = None + try: + self.repl_projects.close() + except BaseException as error: + repl_error = error + try: + self.lsp_projects.close() + except BaseException as lsp_error: + if repl_error is None: + raise + add_note = getattr(repl_error, "add_note", None) + if add_note is not None: + add_note(f"Lean LSP cleanup also failed: {lsp_error}") + else: # pragma: no cover - Python 3.10 compatibility + logger.error("Lean LSP cleanup also failed: %s", lsp_error) + if repl_error is not None: + raise repl_error.with_traceback(repl_error.__traceback__) def _acquisition_timeout(self, operation_timeout: float) -> float: """Reserve enough of the RPC deadline for the admitted tool operation.""" @@ -877,24 +1148,57 @@ class LeanRuntimeRequestHandler(socketserver.StreamRequestHandler): server: LeanRuntimeServer - def setup(self) -> None: - self.request.settimeout(self.server.services.config.rpc_read_timeout) - super().setup() + def _read_request(self) -> bytes: + deadline = ( + time.monotonic() + self.server.services.config.rpc_read_timeout + ) + data = bytearray() + while len(data) <= MAX_MESSAGE_BYTES: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("runtime request read deadline expired") + self.request.settimeout(remaining) + chunk = self.request.recv( + min(65536, MAX_MESSAGE_BYTES + 1 - len(data)) + ) + if time.monotonic() >= deadline: + raise TimeoutError("runtime request read deadline expired") + if not chunk: + raise ValueError("request is empty or unterminated") + data.extend(chunk) + newline = data.find(b"\n") + if newline >= 0: + if len(data) > MAX_MESSAGE_BYTES: + raise ValueError("request exceeds the message limit") + if newline == 0: + raise ValueError("request is empty") + if data[newline + 1 :]: + raise ValueError("request contains trailing data") + return bytes(data[:newline]) + raise ValueError("request exceeds the message limit") def handle(self) -> None: request_id: Any = None + method: Any = None + client_deadline: float | None = None shutdown = False try: - raw = self.rfile.readline(MAX_MESSAGE_BYTES + 1) - if not raw or len(raw) > MAX_MESSAGE_BYTES or not raw.endswith(b"\n"): - raise ValueError("request is empty, unterminated, or too large") - request = json.loads(raw) + raw = self._read_request() + request = _decode_runtime_request(raw) if not isinstance(request, dict): raise ValueError("request must be a JSON object") - request_id = request.get("id") - if request.get("v") != PROTOCOL_VERSION: + candidate_id = request.get("id") + if ( + type(candidate_id) is not str + or not candidate_id + or len(candidate_id) > MAX_REQUEST_ID_CHARS + ): + raise ValueError("id must be a bounded non-empty string") + request_id = candidate_id + version = request.get("v") + if type(version) is not int or version != PROTOCOL_VERSION: raise ValueError( - f"protocol mismatch: expected {PROTOCOL_VERSION}, got {request.get('v')!r}" + f"protocol mismatch: expected {PROTOCOL_VERSION}, got {version!r}" ) method = request.get("method") params = request.get("params") @@ -902,12 +1206,33 @@ def handle(self) -> None: raise ValueError("method must be a non-empty string") if not isinstance(params, dict): raise ValueError("params must be an object") + base_keys = {"v", "id", "method", "params"} + expected_keys = ( + base_keys | {"deadline"} + if method == "repl.run" and "deadline" in request + else base_keys + ) + if set(request) != expected_keys: + raise ValueError("request has an invalid envelope") + candidate_deadline = request.get("deadline") + if "deadline" in request and ( + isinstance(candidate_deadline, bool) + or not isinstance(candidate_deadline, (int, float)) + or not math.isfinite(candidate_deadline) + ): + raise ValueError("deadline must be a finite number or null") + if candidate_deadline is not None: + client_deadline = float(candidate_deadline) if method == "daemon.shutdown": result = {"stopping": True, "pid": os.getpid()} shutdown = True else: - result = self.server.services.dispatch(method, params) + result = self.server.services.dispatch( + method, + params, + client_deadline=client_deadline, + ) response = { "v": PROTOCOL_VERSION, "id": request_id, @@ -928,8 +1253,19 @@ def handle(self) -> None: encoded = json.dumps(response, separators=(",", ":")).encode("utf-8") + b"\n" if len(encoded) > MAX_MESSAGE_BYTES: - encoded = json.dumps( - { + if method == "repl.run" and response.get("ok") is True: + response = { + "v": PROTOCOL_VERSION, + "id": request_id, + "ok": True, + "result": ( + "REPL error (execution outcome unknown; request not retried): " + "the completed response exceeded the runtime message limit; " + "the request must not be replayed." + ), + } + else: + response = { "v": PROTOCOL_VERSION, "id": request_id, "ok": False, @@ -937,18 +1273,53 @@ def handle(self) -> None: "type": "ValueError", "message": "response exceeds the message limit", }, - }, - separators=(",", ":"), - ).encode("utf-8") + b"\n" + } + encoded = json.dumps(response, separators=(",", ":")).encode("utf-8") + b"\n" + write_timeout = min( + getattr( + self.server.services.config, + "response_timeout", + DEFAULT_RESPONSE_TIMEOUT, + ), + RUNTIME_SAFETY_SECONDS, + ) + write_deadline = time.monotonic() + write_timeout + if method == "repl.run" and client_deadline is not None: + write_deadline = min( + write_deadline, + client_deadline + REPL_RESPONSE_GRACE_SECONDS, + ) try: + remaining = write_deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("runtime response write deadline expired") + self.request.settimeout(remaining) self.wfile.write(encoded) self.wfile.flush() - except BrokenPipeError: - logger.warning("Lean runtime client disconnected before receiving its response") + except (TimeoutError, OSError): + logger.warning( + "Lean runtime client disconnected before receiving its response", + exc_info=True, + ) if shutdown: self.server.request_shutdown() +def _close_services_until_clean(services: LeanRuntimeServices) -> None: + """Keep terminal ownership until every quarantined child is gone.""" + delay = TERMINAL_CLEANUP_RETRY_SECONDS + while True: + try: + services.close() + return + except Exception: + logger.exception( + "Lean runtime cleanup remains incomplete; retaining ownership" + ) + time.sleep(delay) + delay = min(delay * 2, MAX_TERMINAL_CLEANUP_RETRY_SECONDS) + + def _configure_logging(log_path: Path | None) -> None: handlers: list[logging.Handler] = [] if log_path is not None: @@ -1014,7 +1385,7 @@ def request_shutdown(signum: int, frame: Any) -> None: finally: try: if services is not None: - services.close() + _close_services_until_clean(services) finally: try: info = paths.socket.lstat() @@ -1029,22 +1400,33 @@ def request_shutdown(signum: int, frame: Any) -> None: os.close(lifetime_fd) -def _paths_from_args(socket_path: str | None, log_path: str | None) -> RuntimePaths: +def _paths_from_args( + socket_path: str | None, + log_path: str | None, + lifetime_lock_path: str | None = None, +) -> RuntimePaths: paths = ( runtime_paths_for_socket(socket_path) if socket_path is not None else default_runtime_paths() ) - if log_path is None: + if log_path is None and lifetime_lock_path is None: return paths - log = Path(log_path).expanduser() + log = paths.log if log_path is None else Path(log_path).expanduser() + lifetime_lock = ( + paths.lifetime_lock + if lifetime_lock_path is None + else Path(lifetime_lock_path).expanduser() + ) if not log.is_absolute(): raise LeanRuntimeError("Lean runtime log path must be absolute") + if not lifetime_lock.is_absolute(): + raise LeanRuntimeError("Lean runtime lifetime lock path must be absolute") return RuntimePaths( directory=paths.directory, socket=paths.socket, lock=paths.lock, - lifetime_lock=paths.lifetime_lock, + lifetime_lock=lifetime_lock, log=log, ) @@ -1053,6 +1435,7 @@ def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--socket", help="override the Unix socket path") parser.add_argument("--log", help="override the rotating log path") + parser.add_argument("--lifetime-lock", help="override the runtime lifetime lock") parser.add_argument( "command", choices=("serve", "start", "status", "stop"), @@ -1060,7 +1443,7 @@ def main(argv: list[str] | None = None) -> None: default="status", ) args = parser.parse_args(argv) - paths = _paths_from_args(args.socket, args.log) + paths = _paths_from_args(args.socket, args.log, args.lifetime_lock) if args.command == "serve": serve(paths) diff --git a/servers/lsp/server.py b/servers/lsp/server.py index 35af19ab..40f09d83 100644 --- a/servers/lsp/server.py +++ b/servers/lsp/server.py @@ -15,9 +15,10 @@ from dataclasses import dataclass, field from logging import getLogger from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any -from fastmcp.server import FastMCP +if TYPE_CHECKING: + from fastmcp.server import FastMCP from servers import resolve_lean_project_dir from servers.lean_client import LeanRuntimeClient @@ -555,6 +556,8 @@ def format_lsp_diagnostics(diagnostics: list[dict]) -> str: def create_lsp_server(runtime: LeanRuntimeClient) -> FastMCP: """Create the public LSP MCP adapter for the shared Lean runtime.""" + from fastmcp.server import FastMCP + server = FastMCP(name="autoform-lsp") @server.tool diff --git a/servers/repl/core.py b/servers/repl/core.py index 7d188017..2f8cd80b 100644 --- a/servers/repl/core.py +++ b/servers/repl/core.py @@ -530,6 +530,14 @@ def __init__(self, message: str, response: Any) -> None: self.response = response +class ReplCleanupError(RuntimeError): + """A disposable call produced a result but its process was not reaped.""" + + def __init__(self, message: str, result: dict[str, Any]) -> None: + super().__init__(message) + self.result = result + + class LeanRepl: """Lean REPL process manager. @@ -566,15 +574,27 @@ def __init__(self, config: LeanReplConfig) -> None: def start( self, startup_timeout: float | None = None, + *, + deadline: float | None = None, + warmup_imports: frozenset[str] | tuple[str, ...] | None = None, ) -> None: """Start and warm the Lean REPL within one startup deadline.""" if os.name != "posix": raise RuntimeError("Lean REPL transport requires a POSIX platform") - timeout = self.config.startup_timeout if startup_timeout is None else min( - self.config.startup_timeout, - startup_timeout, - ) - deadline = time.monotonic() + timeout + if startup_timeout is not None and deadline is not None: + raise TypeError("pass startup_timeout or deadline, not both") + started = time.monotonic() + configured_deadline = started + self.config.startup_timeout + if deadline is None: + timeout = ( + self.config.startup_timeout + if startup_timeout is None + else min(self.config.startup_timeout, startup_timeout) + ) + deadline = started + timeout + else: + deadline = min(deadline, configured_deadline) + timeout = max(0.0, deadline - started) def remaining() -> float: value = deadline - time.monotonic() @@ -586,6 +606,7 @@ def remaining() -> float: env.update(self.config.env) try: + remaining() self.process = subprocess.Popen( self.config.repl_command, cwd=self.cwd, @@ -600,15 +621,20 @@ def remaining() -> float: self._stderr_bytes = 0 self._stderr_tail.clear() - if self.config.warmup_imports: - header = "\n".join( - f"import {root}" for root in self.config.warmup_imports - ) - logger.info( - "Loading imports at startup: %s", - self.config.warmup_imports, + startup_imports = ( + self.config.warmup_imports + if warmup_imports is None + else warmup_imports + ) + if startup_imports: + header = "\n".join(f"import {root}" for root in startup_imports) + logger.info("Loading imports at startup: %s", startup_imports) + resp = self._run( + code=header, + env_id=None, + timeout=remaining(), + deadline=deadline, ) - resp = self._run(code=header, env_id=None, timeout=remaining()) environment, messages = _validate_command_response( resp, context="startup imports", @@ -625,6 +651,7 @@ def remaining() -> float: code="#check Nat", env_id=self._base_env_id, timeout=min(DEFAULT_SMOKE_TEST_TIMEOUT, remaining()), + deadline=deadline, ) _, smoke_messages = _validate_command_response( smoke, @@ -685,6 +712,14 @@ def close(self, *, deadline: float | None = None) -> None: self._stderr_bytes = 0 self._stderr_tail.clear() + def close_with_deadline(self, deadline: float) -> None: + """Close using an absolute deadline shared by a pool shutdown.""" + self.close(deadline=deadline) + + def is_clean(self) -> bool: + """Return whether this wrapper owns no live or unreaped process group.""" + return self.process is None and self._process_group_id is None + def restart(self, timeout: float | None = None) -> None: """Restart the Lean REPL process within an optional total timeout.""" deadline = time.monotonic() + timeout if timeout is not None else None @@ -713,6 +748,133 @@ def get_memory_usage(self) -> float: """Return memory usage in GB.""" return _get_process_memory_gb(self.process) + def run_disposable( + self, + code: str, + timeout: float | None = None, + *, + deadline: float | None = None, + ) -> dict[str, Any]: + """Run one public call as the only frame sent to a fresh process.""" + if timeout is not None and deadline is not None: + raise TypeError("pass timeout or deadline, not both") + if deadline is None: + timeout = self.request_timeout if timeout is None else timeout + deadline = time.monotonic() + timeout + else: + timeout = max(0.0, deadline - time.monotonic()) + + def remaining() -> float: + value = deadline - time.monotonic() + if value <= 0: + raise TimeoutError(f"REPL command timed out after {timeout:g} seconds") + return value + + with self._process_lock: + result: dict[str, Any] | None = None + request_error: BaseException | None = None + try: + self.close(deadline=deadline) + imports, _, _ = _split_imports_and_body(code) + if ( + self.config.validate_imports + and self._allowed_import_roots is not None + ): + submitted_roots = { + statement.split(".")[0] for statement in imports + } + disallowed = submitted_roots - self._allowed_import_roots + if disallowed: + result = { + "repl_error": ( + f"Disallowed imports: {', '.join(sorted(disallowed))}. " + "Allowed roots: " + f"{', '.join(sorted(self._allowed_import_roots))}." + ) + } + if result is None: + added_imports = tuple( + root + for root in sorted(self.config.warmup_imports) + if root not in imports + ) + prefix = "\n".join(f"import {root}" for root in added_imports) + command = f"{prefix}\n{code}" if prefix else code + self.start(deadline=deadline, warmup_imports=()) + response = self._run( + code=command, + env_id=None, + timeout=remaining(), + deadline=deadline, + ) + _validate_command_response( + response, + context="the requested command", + require_environment=True, + ) + _adjust_line_numbers(response, -len(added_imports)) + result = _without_process_handles(response) + except ReplStderrBacklog as error: + try: + _validate_command_response( + error.response, + context="the requested command", + require_environment=True, + ) + except ReplCommandError as command_error: + result = {"repl_error": str(command_error)} + except ReplProtocolError as protocol_error: + result = { + "repl_error": str(protocol_error), + "outcome_unknown": True, + } + else: + response = _without_process_handles(error.response) + _adjust_line_numbers(response, -len(added_imports)) + result = response + except ReplCommandError as error: + result = {"repl_error": str(error)} + except ReplProtocolError as error: + result = {"repl_error": str(error), "outcome_unknown": True} + except ReplOutcomeUnknown as error: + result = {"repl_error": str(error), "outcome_unknown": True} + except (ReplProcessExited, TimeoutError, RuntimeError) as error: + result = {"repl_error": str(error)} + except BaseException as error: + request_error = error + finally: + try: + self.close( + deadline=time.monotonic() + DEFAULT_REPL_CLEANUP_SECONDS + ) + except BaseException as cleanup_error: + logger.exception( + "failed to retire disposable Lean REPL process; " + "the worker must not be reused" + ) + if request_error is not None: + note = f"Lean REPL process cleanup also failed: {cleanup_error}" + add_note = getattr(request_error, "add_note", None) + if add_note is not None: + add_note(note) + else: # pragma: no cover - Python 3.10 compatibility + logger.error("%s", note) + elif not isinstance(cleanup_error, Exception): + raise + elif result is not None: + raise ReplCleanupError( + "Disposable Lean REPL process cleanup failed after a " + "result was produced; the result must not be replayed: " + f"{cleanup_error}", + result, + ) from cleanup_error + + if request_error is not None: + raise request_error.with_traceback(request_error.__traceback__) + if result is None: + raise RuntimeError("disposable Lean REPL call produced no result") + return result + def run(self, code: str, env_id: int | None = None, timeout: float | None = None) -> dict[str, Any]: """Send code to the REPL within one deadline across recovery attempts.""" timeout = self.request_timeout if timeout is None else timeout @@ -886,17 +1048,36 @@ def _check_memory_and_maybe_restart(self, timeout: float | None = None) -> None: except Exception: logger.warning("Memory check failed, continuing", exc_info=True) - def _run(self, code: str, env_id: int | None, timeout: float) -> dict[str, Any]: + def _run( + self, + code: str, + env_id: int | None, + timeout: float, + *, + deadline: float | None = None, + ) -> dict[str, Any]: """Run one frame and distinguish safe pre-send failures from unknown outcomes.""" request_sent = False - cleanup_deadline = time.monotonic() + timeout + started = time.monotonic() + cleanup_deadline = started + timeout + if deadline is not None: + cleanup_deadline = min(cleanup_deadline, deadline) + remaining = cleanup_deadline - started + if remaining <= 0: + raise TimeoutError(f"REPL command timed out after {timeout:g} seconds") def mark_sent() -> None: nonlocal request_sent request_sent = True try: - return self._run_io(code, env_id, timeout, mark_sent) + return self._run_io( + code, + env_id, + remaining, + mark_sent, + deadline=cleanup_deadline, + ) except ReplOutcomeUnknown as error: message = str(error) try: @@ -939,6 +1120,8 @@ def _run_io( env_id: int | None, timeout: float, mark_sent: Callable[[], None], + *, + deadline: float | None = None, ) -> dict[str, Any]: """Send code to the REPL via stdin JSON-RPC, read response via non-blocking I/O.""" cmd_obj: dict[str, Any] = {"cmd": code} @@ -955,7 +1138,7 @@ def _run_io( ): raise ReplProcessExited("REPL process is not running.") - end_time = time.monotonic() + timeout + end_time = time.monotonic() + timeout if deadline is None else deadline stdin_fd = self.process.stdin.fileno() stdout_fd = self.process.stdout.fileno() stderr_fd = self.process.stderr.fileno() diff --git a/servers/repl/pool.py b/servers/repl/pool.py index ff543a5f..8de493c5 100644 --- a/servers/repl/pool.py +++ b/servers/repl/pool.py @@ -9,13 +9,27 @@ from logging import getLogger from typing import Any -from .core import LeanRepl, LeanReplConfig +from .core import ( + DEFAULT_REPL_CLEANUP_SECONDS, + LeanRepl, + LeanReplConfig, + ReplCleanupError, +) logger = getLogger(__name__) DEFAULT_PORT = 8990 DEFAULT_RAM_FRACTION = 0.5 DEFAULT_STARTUP_STAGGER_SECONDS = 2.0 +DEFAULT_POOL_CLEANUP_SECONDS = DEFAULT_REPL_CLEANUP_SECONDS + + +class ReplPoolBusyError(TimeoutError): + """A REPL request expired before any worker could receive it.""" + + +class ReplPoolUnavailableError(RuntimeError): + """A REPL pool stopped before any worker could receive the request.""" @dataclass @@ -37,10 +51,10 @@ def __post_init__(self) -> None: class LeanReplPool: - """Pool of Lean REPL instances with queue-based load balancing. + """Pool of cold Lean REPL slots with queue-based load balancing. - Each worker thread owns its own LeanRepl subprocess. Tasks are - distributed to idle workers via a FIFO queue. + Each slot owns a ``LeanRepl`` wrapper, but no subprocess survives a public + request. Tasks are distributed to idle slots via a FIFO queue. """ def __init__(self, config: LeanReplPoolConfig) -> None: @@ -51,77 +65,207 @@ def __init__(self, config: LeanReplPoolConfig) -> None: self._workers: list[LeanRepl] = [] self._idle: queue.Queue[LeanRepl] = queue.Queue() self._lock = threading.Lock() + self._condition = threading.Condition(self._lock) + self._active_calls = 0 + self._closing = False + self._closed = False try: - for i in range(self.capacity): - if i > 0: - import time - - time.sleep(config.startup_stagger) + for _ in range(self.capacity): repl = LeanRepl(config) - try: - repl.start() - except BaseException: - # LeanRepl.start() currently cleans up its own process, but - # keep the pool transaction safe for alternate/test workers - # and future implementations too. - try: - repl.close() - except Exception: - logger.exception("failed to close REPL after startup error") - raise self._workers.append(repl) self._idle.put(repl) except BaseException: - self._close_workers() + try: + self.shutdown() + except BaseException: + logger.exception("failed to clean up a partially constructed REPL pool") raise - def _close_workers(self) -> None: - """Close every constructed worker, preserving cleanup after one failure.""" - for worker in reversed(self._workers): + @staticmethod + def _close_worker(worker: LeanRepl, deadline: float | None = None) -> None: + close_with_deadline = getattr(worker, "close_with_deadline", None) + if deadline is not None and close_with_deadline is not None: + close_with_deadline(deadline) + return + worker.close() + + def _close_workers(self, deadline: float) -> None: + """Close every worker and retain any whose cleanup failed.""" + failed_workers = [] + first_error: BaseException | None = None + for worker in self._workers: try: - worker.close() - except Exception: + self._close_worker(worker, deadline) + except BaseException as error: logger.exception("failed to close REPL worker") - self._workers.clear() + failed_workers.append(worker) + if first_error is None: + first_error = error + self._workers = failed_workers while True: try: self._idle.get_nowait() except queue.Empty: break + if first_error is not None: + raise first_error def run(self, code: str, **kwargs: Any) -> dict[str, Any]: """Run code on an idle REPL within one queue-and-execution timeout.""" timeout = kwargs.pop("timeout", None) - deadline = time.monotonic() + timeout if timeout is not None else None - try: - repl = self._idle.get(timeout=timeout) - except queue.Empty as error: - raise TimeoutError( - f"timed out after {timeout:g}s waiting for an idle Lean REPL" - ) from error + deadline = kwargs.pop("deadline", None) + if kwargs: + names = ", ".join(sorted(kwargs)) + raise TypeError(f"unsupported Lean REPL pool arguments: {names}") + if timeout is not None and deadline is not None: + raise TypeError("pass timeout or deadline, not both") + if deadline is None and timeout is not None: + deadline = time.monotonic() + timeout + timeout_description = f" after {timeout:g}s" if timeout is not None else "" + with self._condition: + if self._shutdown: + raise ReplPoolUnavailableError("Lean REPL pool is shut down") + self._active_calls += 1 + repl: LeanRepl | None = None def run_once() -> dict[str, Any]: - call_kwargs = dict(kwargs) if deadline is not None: remaining = deadline - time.monotonic() if remaining <= 0: - raise TimeoutError( - f"timed out after {timeout:g}s waiting for an idle Lean REPL" + raise ReplPoolBusyError( + f"timed out{timeout_description} waiting for an idle Lean REPL" ) - call_kwargs["timeout"] = remaining - return repl.run(code, **call_kwargs) + return repl.run_disposable(code, deadline=deadline) + return repl.run_disposable(code) + result: dict[str, Any] | None = None + request_error: BaseException | None = None try: - return run_once() + while repl is None: + with self._condition: + if self._shutdown: + raise ReplPoolUnavailableError("Lean REPL pool is shut down") + wait = 0.1 + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise ReplPoolBusyError( + f"timed out{timeout_description} waiting for an idle Lean REPL" + ) + wait = min(wait, remaining) + try: + repl = self._idle.get(timeout=wait) + except queue.Empty: + continue + with self._condition: + if self._shutdown: + raise ReplPoolUnavailableError("Lean REPL pool is shut down") + result = run_once() + except BaseException as error: + request_error = error finally: - self._idle.put(repl) + reusable = repl is None + cleanup_error: BaseException | None = None + if repl is not None: + try: + self._close_worker( + repl, + time.monotonic() + DEFAULT_POOL_CLEANUP_SECONDS, + ) + reusable = getattr(repl, "is_clean", lambda: True)() + except BaseException as error: + cleanup_error = error + finally: + if not reusable and cleanup_error is None: + cleanup_error = RuntimeError( + "Lean REPL process cleanup could not be confirmed" + ) + with self._condition: + if not reusable: + self._shutdown = True + if reusable and not self._shutdown: + self._idle.put(repl) + self._active_calls -= 1 + self._condition.notify_all() + else: + with self._condition: + self._active_calls -= 1 + self._condition.notify_all() + + if isinstance(request_error, ReplCleanupError): + if cleanup_error is None: + return request_error.result + if not isinstance(cleanup_error, Exception): + raise cleanup_error.with_traceback(cleanup_error.__traceback__) + result = request_error.result + if result.get("outcome_unknown") is True: + result = dict(result) + result["repl_error"] = ( + f"{result['repl_error']}; process cleanup also failed: " + f"{cleanup_error}" + ) + return result + return { + "repl_error": ( + "Lean REPL command may have completed, but process cleanup " + f"could not be confirmed: {cleanup_error}. The request was not " + "retried and must not be replayed." + ), + "outcome_unknown": True, + } + if request_error is not None: + if cleanup_error is not None: + note = f"Lean REPL process cleanup also failed: {cleanup_error}" + add_note = getattr(request_error, "add_note", None) + if add_note is not None: + add_note(note) + else: # pragma: no cover - Python 3.10 compatibility + logger.error("%s", note) + raise request_error.with_traceback(request_error.__traceback__) + if cleanup_error is not None and not isinstance(cleanup_error, Exception): + raise cleanup_error.with_traceback(cleanup_error.__traceback__) + if result is None: + if cleanup_error is not None: + raise cleanup_error.with_traceback(cleanup_error.__traceback__) + raise RuntimeError("Lean REPL pool call produced no result") + if cleanup_error is not None: + return { + "repl_error": ( + "Lean REPL command may have completed, but process cleanup " + f"could not be confirmed: {cleanup_error}. The request was not " + "retried and must not be replayed." + ), + "outcome_unknown": True, + } + return result def get_memory_usage(self) -> float: """Total memory usage across all REPL instances in GB.""" return sum(w.get_memory_usage() for w in self._workers) + def is_usable(self) -> bool: + """Return whether the pool can admit another request.""" + with self._condition: + return not self._shutdown and not self._closed + def shutdown(self) -> None: """Shut down all REPL instances.""" - self._shutdown = True - self._close_workers() + with self._condition: + self._shutdown = True + self._condition.notify_all() + while self._active_calls: + self._condition.wait() + while self._closing: + self._condition.wait() + if self._closed: + return + self._closing = True + try: + deadline = time.monotonic() + DEFAULT_POOL_CLEANUP_SECONDS + self._close_workers(deadline) + finally: + with self._condition: + self._closing = False + self._closed = not self._workers + self._condition.notify_all() diff --git a/servers/repl/server.py b/servers/repl/server.py index d0a72c3d..656a5d00 100644 --- a/servers/repl/server.py +++ b/servers/repl/server.py @@ -3,19 +3,23 @@ from __future__ import annotations import json +from typing import TYPE_CHECKING -from fastmcp.server import FastMCP +if TYPE_CHECKING: + from fastmcp.server import FastMCP from servers.lean_client import LeanRuntimeClient def create_repl_server(runtime: LeanRuntimeClient) -> FastMCP: """Create the public REPL MCP adapter for the shared Lean runtime.""" + from fastmcp.server import FastMCP + server = FastMCP(name="autoform-repl") @server.tool def run_lean_code(project_dir: str, code: str, timeout: float | None = None) -> str: - """Compile a Lean snippet in a project's persistent REPL. + """Compile a Lean snippet in a fresh project-scoped REPL process. Args: project_dir: Absolute path to the Lake project root. diff --git a/tests/fixtures/repl-smoke/.gitignore b/tests/fixtures/repl-smoke/.gitignore new file mode 100644 index 00000000..01f8cdb6 --- /dev/null +++ b/tests/fixtures/repl-smoke/.gitignore @@ -0,0 +1 @@ +.lake/ diff --git a/tests/fixtures/repl-smoke/lake-manifest.json b/tests/fixtures/repl-smoke/lake-manifest.json new file mode 100644 index 00000000..59becf30 --- /dev/null +++ b/tests/fixtures/repl-smoke/lake-manifest.json @@ -0,0 +1,16 @@ +{"version": "1.2.0", + "packagesDir": ".lake/packages", + "packages": + [{"url": "https://github.com/leanprover-community/repl.git", + "type": "git", + "subDir": null, + "scope": "", + "rev": "68a3b3a059787a7db44fb1e6281e4a657efee470", + "name": "repl", + "manifestFile": "lake-manifest.json", + "inputRev": "68a3b3a059787a7db44fb1e6281e4a657efee470", + "inherited": false, + "configFile": "lakefile.toml"}], + "name": "AutoformReplSmoke", + "lakeDir": ".lake", + "fixedToolchain": false} diff --git a/tests/fixtures/repl-smoke/lakefile.toml b/tests/fixtures/repl-smoke/lakefile.toml new file mode 100644 index 00000000..6e845fe5 --- /dev/null +++ b/tests/fixtures/repl-smoke/lakefile.toml @@ -0,0 +1,7 @@ +name = "AutoformReplSmoke" +version = "0.1.0" + +[[require]] +name = "repl" +git = "https://github.com/leanprover-community/repl.git" +rev = "68a3b3a059787a7db44fb1e6281e4a657efee470" diff --git a/tests/fixtures/repl-smoke/lean-toolchain b/tests/fixtures/repl-smoke/lean-toolchain new file mode 100644 index 00000000..94b9f495 --- /dev/null +++ b/tests/fixtures/repl-smoke/lean-toolchain @@ -0,0 +1 @@ +leanprover/lean4:v4.32.0 diff --git a/tests/test_real_repl.py b/tests/test_real_repl.py new file mode 100644 index 00000000..81dc946d --- /dev/null +++ b/tests/test_real_repl.py @@ -0,0 +1,38 @@ +"""Opt-in integration test against the pinned upstream Lean REPL.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from servers.repl.core import LeanRepl, LeanReplConfig + + +REPL_FIXTURE = Path(__file__).parent / "fixtures" / "repl-smoke" + + +@pytest.mark.skipif( + os.environ.get("AUTOFORM_RUN_REAL_REPL_TESTS") != "1", + reason="set AUTOFORM_RUN_REAL_REPL_TESTS=1 to run the pinned REPL integration", +) +def test_disposable_call_matches_the_pinned_repl_protocol(): + repl = LeanRepl( + LeanReplConfig( + cwd=str(REPL_FIXTURE), + repl_command=["lake", "exe", "repl"], + warmup_imports=frozenset(), + validate_imports=False, + ) + ) + + response = repl.run_disposable( + "theorem autoform_repl_probe : True := by sorry", + timeout=180, + ) + + assert response.get("sorries") + assert "env" not in response + assert all("proofState" not in sorry for sorry in response["sorries"]) + assert repl.is_clean() diff --git a/tests/test_repl_core_protocol.py b/tests/test_repl_core_protocol.py index 06622e4d..09ca1637 100644 --- a/tests/test_repl_core_protocol.py +++ b/tests/test_repl_core_protocol.py @@ -49,6 +49,33 @@ def popen(*args, **kwargs): repl._process_group_id = None +def test_start_does_not_spawn_after_environment_setup_exhausts_deadline(monkeypatch): + now = [100.0] + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + validate_imports=False, + warmup_imports=frozenset(), + ) + ) + + def clean_environment(): + now[0] = 102.0 + return {} + + monkeypatch.setattr(repl_core.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(repl_core, "_inherit_clean_env", clean_environment) + monkeypatch.setattr( + repl_core.subprocess, + "Popen", + lambda *args, **kwargs: pytest.fail( + "an expired startup deadline must not spawn a Lean process" + ), + ) + + with pytest.raises(TimeoutError, match="startup timed out"): + repl.start(startup_timeout=1) + + def test_start_rejects_unsupported_platform_before_spawning(monkeypatch): monkeypatch.setattr(repl_core.os, "name", "nt") monkeypatch.setattr( @@ -99,8 +126,7 @@ def pid(self): repl.start() assert retired == [(process, 4321)] - assert repl.process is None - assert repl._process_group_id is None + assert repl.is_clean() is True @pytest.mark.skipif(os.name != "posix", reason="process groups require POSIX") @@ -669,6 +695,217 @@ def test_additive_response_fields_are_tolerated_but_not_exported(): } +def test_run_disposable_sends_one_frame_and_removes_process_handles(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + validate_imports=False, + warmup_imports=frozenset({"Mathlib"}), + ) + ) + repl.process = object() + events = [] + + def close(*, deadline=None): + events.append(("close", repl._process_lock.locked())) + repl.process = None + + def start(startup_timeout=None, *, deadline=None, warmup_imports=None): + events.append(("start", startup_timeout, deadline, warmup_imports)) + repl.process = object() + + def run_frame(code, env_id, timeout, *, deadline=None): + events.append(("frame", code, env_id, timeout, deadline, repl.process)) + return { + "env": 7, + "messages": [], + "sorries": [{"goal": "False", "proofState": 9}], + } + + monkeypatch.setattr(repl, "close", close) + monkeypatch.setattr(repl, "start", start) + monkeypatch.setattr(repl, "_run", run_frame) + + response = repl.run_disposable("#check Nat", timeout=3) + + assert events[0] == ("close", True) + assert events[1][0] == "start" + assert events[1][1] is None + assert events[1][2] is not None + assert events[1][3] == () + assert events[2][0:3] == ("frame", "import Mathlib\n#check Nat", None) + assert 0 < events[2][3] <= 3 + assert events[2][4] == events[1][2] + assert events[2][5] is not None + assert events[3] == ("close", True) + assert len([event for event in events if event[0] == "frame"]) == 1 + assert response == {"messages": [], "sorries": [{"goal": "False"}]} + assert repl.process is None + + +def test_run_disposable_reserves_cleanup_time_after_command_deadline(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig(validate_imports=False, warmup_imports=frozenset()) + ) + now = [100.0] + close_deadlines = [] + + def close(*, deadline=None): + close_deadlines.append(deadline) + + def run_frame(*args, **kwargs): + now[0] = 102.999 + return {"env": 1, "messages": [], "sorries": []} + + monkeypatch.setattr(repl_core.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(repl, "close", close) + monkeypatch.setattr(repl, "start", lambda *args, **kwargs: None) + monkeypatch.setattr(repl, "_run", run_frame) + + assert repl.run_disposable("#check Nat", timeout=3) == { + "messages": [], + "sorries": [], + } + assert close_deadlines == [ + 103.0, + 102.999 + repl_core.DEFAULT_REPL_CLEANUP_SECONDS, + ] + + +def test_run_disposable_closes_after_frame_error(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig(validate_imports=False, warmup_imports=frozenset()) + ) + close_calls = [] + monkeypatch.setattr( + repl, + "close", + lambda *, deadline=None: close_calls.append(True), + ) + monkeypatch.setattr(repl, "start", lambda *args, **kwargs: None) + monkeypatch.setattr( + repl, + "_run", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("failed")), + ) + + assert repl.run_disposable("#check Nat") == {"repl_error": "failed"} + + assert close_calls == [True, True] + + +def test_run_disposable_does_not_swallow_cleanup_cancellation(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig(validate_imports=False, warmup_imports=frozenset()) + ) + close_calls = 0 + + def close(*, deadline=None): + nonlocal close_calls + close_calls += 1 + if close_calls == 2: + raise asyncio.CancelledError + + monkeypatch.setattr(repl, "close", close) + monkeypatch.setattr(repl, "start", lambda *args, **kwargs: None) + monkeypatch.setattr( + repl, + "_run", + lambda *args, **kwargs: {"env": 1, "messages": [], "sorries": []}, + ) + + with pytest.raises(asyncio.CancelledError): + repl.run_disposable("#check Nat") + + assert close_calls == 2 + + +def test_run_disposable_preserves_request_cancellation_when_cleanup_also_cancels( + monkeypatch, +): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig(validate_imports=False, warmup_imports=frozenset()) + ) + close_calls = 0 + + def close(*, deadline=None): + nonlocal close_calls + close_calls += 1 + if close_calls == 2: + raise asyncio.CancelledError("cleanup") + + monkeypatch.setattr(repl, "close", close) + monkeypatch.setattr(repl, "start", lambda *args, **kwargs: None) + monkeypatch.setattr( + repl, + "_run", + lambda *args, **kwargs: (_ for _ in ()).throw(KeyboardInterrupt("request")), + ) + + with pytest.raises(KeyboardInterrupt, match="request") as raised: + repl.run_disposable("#check Nat") + + if hasattr(raised.value, "add_note"): + assert raised.value.__notes__ == [ + "Lean REPL process cleanup also failed: cleanup" + ] + assert close_calls == 2 + + +def test_run_disposable_never_returns_success_before_verified_cleanup(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig(validate_imports=False, warmup_imports=frozenset()) + ) + close_calls = 0 + + def close(*, deadline=None): + nonlocal close_calls + close_calls += 1 + if close_calls == 2: + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(repl, "close", close) + monkeypatch.setattr(repl, "start", lambda *args, **kwargs: None) + monkeypatch.setattr( + repl, + "_run", + lambda *args, **kwargs: {"env": 1, "messages": [], "sorries": []}, + ) + + with pytest.raises(repl_core.ReplCleanupError) as raised: + repl.run_disposable("#check Nat") + + assert raised.value.result == {"messages": [], "sorries": []} + assert close_calls == 2 + + +def test_run_disposable_closes_before_rejecting_an_import(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + allowed_imports=frozenset({"Mathlib"}), + warmup_imports=frozenset(), + ) + ) + repl.process = object() + close_calls = [] + + def close(*, deadline=None): + close_calls.append(True) + repl.process = None + + monkeypatch.setattr(repl, "close", close) + monkeypatch.setattr( + repl, + "start", + lambda *args, **kwargs: pytest.fail("invalid input must not start Lean"), + ) + + response = repl.run_disposable("import Unsafe\n#check Nat") + + assert "Disallowed imports: Unsafe" in response["repl_error"] + assert close_calls == [True, True] + assert repl.process is None + + class _PipeProcess: def __init__(self, stack: ExitStack, stdout_chunks: list[bytes], stderr: bytes = b""): stdin_read, stdin_write = os.pipe() @@ -737,6 +974,28 @@ def fake_select(readable, writable, exceptional, timeout=None): monkeypatch.setattr(repl_core.select, "select", fake_select) +def test_run_forwards_absolute_deadline_to_wire(monkeypatch): + now = [100.0] + observed = [] + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + warmup_imports=frozenset(), + validate_imports=False, + ) + ) + + def run_io(code, env_id, timeout, mark_sent, *, deadline=None): + now[0] = 104.0 + observed.append((timeout, deadline)) + return {"env": 1, "messages": [], "sorries": []} + + monkeypatch.setattr(repl_core.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(repl, "_run_io", run_io) + + assert repl._run("#check Nat", None, 5, deadline=105.0)["env"] == 1 + assert observed == [(5.0, 105.0)] + + def test_response_timeout_after_full_write_is_not_retried(monkeypatch): repl = repl_core.LeanRepl( repl_core.LeanReplConfig( @@ -760,7 +1019,7 @@ def test_response_timeout_after_full_write_is_not_retried(monkeypatch): lambda timeout=None: pytest.fail("a sent request must not be retried"), ) - def fail_after_send(code, env_id, timeout, mark_sent): + def fail_after_send(code, env_id, timeout, mark_sent, *, deadline=None): calls.append((code, env_id)) mark_sent() raise TimeoutError("response timed out") @@ -797,7 +1056,7 @@ def close(**kwargs): close_calls += 1 raise RuntimeError("cleanup failed") - def fail_after_send(code, env_id, timeout, mark_sent): + def fail_after_send(code, env_id, timeout, mark_sent, *, deadline=None): mark_sent() raise TimeoutError("response timed out") @@ -820,7 +1079,7 @@ def test_run_closes_and_reraises_cancellation(monkeypatch, request_sent): repl.process = object() retired = [] - def cancel(code, env_id, timeout, mark_sent): + def cancel(code, env_id, timeout, mark_sent, *, deadline=None): if request_sent: mark_sent() raise asyncio.CancelledError @@ -1667,9 +1926,7 @@ def slow_reap(process, process_group_id, deadline=None): assert response.get("outcome_unknown") is True assert "cleanup also failed" not in response["repl_error"] - assert repl.process is None - assert repl._process_group_id is None - assert repl._retire_pending is False + assert repl.is_clean() finally: repl.close() diff --git a/tests/test_repl_pool_lifecycle.py b/tests/test_repl_pool_lifecycle.py index 81f36a62..1baf17dd 100644 --- a/tests/test_repl_pool_lifecycle.py +++ b/tests/test_repl_pool_lifecycle.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import threading import pytest @@ -10,20 +11,17 @@ from servers.repl import pool as repl_pool -def test_partial_pool_startup_closes_all_constructed_workers(monkeypatch): +def test_pool_construction_keeps_every_worker_cold(monkeypatch): workers = [] class FakeRepl: def __init__(self, config): - self.number = len(workers) + 1 self.started = False self.closed = False workers.append(self) def start(self): self.started = True - if self.number == 2: - raise RuntimeError("second worker failed") def close(self): self.closed = True @@ -31,13 +29,13 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) config = repl_pool.LeanReplPoolConfig(num_repls=3, startup_stagger=0) - with pytest.raises(RuntimeError, match="second worker failed"): - repl_pool.LeanReplPool(config) + pool = repl_pool.LeanReplPool(config) - assert len(workers) == 2 - assert workers[0].started is True - assert workers[0].closed is True - assert workers[1].closed is True + assert len(workers) == 3 + assert all(worker.started is False for worker in workers) + assert all(worker.closed is False for worker in workers) + assert pool._idle.qsize() == 3 + pool.shutdown() def test_shutdown_closes_every_worker_and_drains_idle_queue(monkeypatch): @@ -66,6 +64,30 @@ def close(self): assert pool._idle.empty() +def test_shutdown_shares_one_cleanup_deadline_across_workers(monkeypatch): + deadlines = [] + + class FakeRepl: + def __init__(self, config): + pass + + def close(self): + pytest.fail("deadline-aware cleanup should be used") + + def close_with_deadline(self, deadline): + deadlines.append(deadline) + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=3, startup_stagger=0) + ) + + pool.shutdown() + + assert len(deadlines) == 3 + assert len(set(deadlines)) == 1 + + def test_request_timeout_includes_waiting_for_an_idle_worker(monkeypatch): class FakeRepl: def __init__(self, config): @@ -90,6 +112,352 @@ def close(self): pool.shutdown() +def test_pool_closes_a_worker_before_requeue_after_request_exception(monkeypatch): + workers = [] + + class FakeRepl: + def __init__(self, config): + self.close_calls = 0 + workers.append(self) + + def run_disposable(self, code, **kwargs): + raise OSError("stdout failed") + + def close(self): + self.close_calls += 1 + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + try: + with pytest.raises(OSError, match="stdout failed"): + pool.run("#check Nat") + + assert workers[0].close_calls == 1 + assert pool._idle.qsize() == 1 + finally: + pool.shutdown() + + +def test_pool_closes_a_worker_before_requeue_after_success(monkeypatch): + events = [] + + class FakeRepl: + def __init__(self, config): + pass + + def run_disposable(self, code, **kwargs): + events.append(("run", code, kwargs)) + return {"messages": []} + + def close(self): + events.append(("close",)) + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + try: + assert pool.run("#check Nat") == {"messages": []} + assert pool._idle.qsize() == 1 + finally: + pool.shutdown() + + assert events[:2] == [("run", "#check Nat", {}), ("close",)] + + +def test_pool_reserves_cleanup_time_after_request_deadline(monkeypatch): + now = [100.0] + close_deadlines = [] + + class FakeRepl: + def __init__(self, config): + pass + + def run_disposable(self, code, **kwargs): + now[0] = 102.999 + return {"messages": []} + + def close_with_deadline(self, deadline): + close_deadlines.append(deadline) + + monkeypatch.setattr(repl_pool.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + try: + assert pool.run("#check Nat", timeout=3) == {"messages": []} + assert close_deadlines == [ + 102.999 + repl_pool.DEFAULT_POOL_CLEANUP_SECONDS + ] + finally: + pool.shutdown() + + +def test_pool_forwards_absolute_deadline_without_resetting_it(monkeypatch): + now = [100.0] + observed = [] + + class FakeRepl: + def __init__(self, config): + pass + + def run_disposable(self, code, **kwargs): + now[0] = 104.0 + observed.append(kwargs) + return {"messages": []} + + def close(self): + pass + + monkeypatch.setattr(repl_pool.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + try: + assert pool.run("#check Nat", deadline=105.0) == {"messages": []} + finally: + pool.shutdown() + + assert observed == [{"deadline": 105.0}] + + +def test_pool_never_requeues_a_worker_that_failed_to_close(monkeypatch): + workers = [] + + class FakeRepl: + def __init__(self, config): + self.close_calls = 0 + workers.append(self) + + def run_disposable(self, code, **kwargs): + raise repl_core.ReplCleanupError( + "first cleanup failed", + {"messages": []}, + ) + + def close(self): + self.close_calls += 1 + if self.close_calls == 1: + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + response = pool.run("#check Nat") + + assert workers[0].close_calls == 1 + assert response["outcome_unknown"] is True + assert "must not be replayed" in response["repl_error"] + assert pool._idle.empty() + assert pool._active_calls == 0 + assert pool._shutdown is True + with pytest.raises(RuntimeError, match="pool is shut down"): + pool.run("#check Bool") + pool.shutdown() + + +def test_pool_returns_captured_response_after_cleanup_retry_succeeds(monkeypatch): + workers = [] + + class FakeRepl: + def __init__(self, config): + self.dirty = True + self.close_calls = 0 + workers.append(self) + + def run_disposable(self, code, **kwargs): + return {"messages": []} + + def close(self): + self.close_calls += 1 + self.dirty = False + + def is_clean(self): + return not self.dirty + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + try: + assert pool.run("#check Nat") == {"messages": []} + assert workers[0].close_calls == 1 + assert pool._idle.qsize() == 1 + assert pool.is_usable() is True + finally: + pool.shutdown() + + +def test_pool_preserves_cancellation_when_cleanup_also_fails(monkeypatch): + class FakeRepl: + def __init__(self, config): + pass + + def run_disposable(self, code, **kwargs): + raise KeyboardInterrupt + + def close(self): + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + with pytest.raises(KeyboardInterrupt) as raised: + pool.run("#check Nat") + + if hasattr(raised.value, "add_note"): + assert raised.value.__notes__ == [ + "Lean REPL process cleanup also failed: cleanup failed" + ] + assert pool._shutdown is True + assert pool._idle.empty() + + +def test_shutdown_retains_and_retries_a_worker_that_failed_to_close(monkeypatch): + workers = [] + + class FakeRepl: + def __init__(self, config): + self.close_calls = 0 + workers.append(self) + + def close(self): + self.close_calls += 1 + if self.close_calls == 1: + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + with pytest.raises(RuntimeError, match="cleanup failed"): + pool.shutdown() + + assert pool._workers == workers + assert pool._closed is False + assert pool._idle.empty() + + pool.shutdown() + + assert workers[0].close_calls == 2 + assert pool._workers == [] + assert pool._closed is True + + +def test_shutdown_never_requeues_a_borrowed_worker(monkeypatch): + running = threading.Event() + release = threading.Event() + shutdown_done = threading.Event() + calls = [] + + class FakeRepl: + def __init__(self, config): + pass + + def run_disposable(self, code, **kwargs): + calls.append(code) + running.set() + release.wait(timeout=2) + return {"messages": []} + + def close(self): + release.set() + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + first = threading.Thread(target=pool.run, args=("first",)) + first.start() + assert running.wait(timeout=1) + + errors = [] + + def wait_for_worker(): + try: + pool.run("second", timeout=0.1) + except (RuntimeError, TimeoutError) as error: + errors.append(error) + + second = threading.Thread(target=wait_for_worker) + second.start() + + def shut_down(): + pool.shutdown() + shutdown_done.set() + + shutdown = threading.Thread(target=shut_down) + shutdown.start() + with pool._condition: + assert pool._condition.wait_for(lambda: pool._shutdown, timeout=1) + assert not shutdown_done.is_set() + release.set() + first.join(timeout=2) + second.join(timeout=2) + shutdown.join(timeout=2) + + assert not first.is_alive() + assert not second.is_alive() + assert not shutdown.is_alive() + assert shutdown_done.is_set() + assert calls == ["first"] + assert len(errors) == 1 + assert pool._idle.empty() + + +def test_concurrent_shutdown_closes_each_worker_once(monkeypatch): + close_started = threading.Event() + release_close = threading.Event() + second_started = threading.Event() + close_calls = [] + + class FakeRepl: + def __init__(self, config): + pass + + def close(self): + close_calls.append(self) + close_started.set() + release_close.wait(timeout=2) + + monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) + pool = repl_pool.LeanReplPool( + repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + ) + + first = threading.Thread(target=pool.shutdown) + + def shut_down_second(): + second_started.set() + pool.shutdown() + + second = threading.Thread(target=shut_down_second) + first.start() + assert close_started.wait(timeout=1) + second.start() + assert second_started.wait(timeout=1) + release_close.set() + first.join(timeout=2) + second.join(timeout=2) + + assert not first.is_alive() + assert not second.is_alive() + assert len(close_calls) == 1 + + def test_repl_retry_recovery_uses_the_original_deadline(monkeypatch): clock = {"now": 0.0} repl = repl_core.LeanRepl( @@ -142,8 +510,6 @@ def poll(self): process = StalledProcess() process.stdin = stdin - # These streams are checked before the bounded write but never read in - # this test because the deliberately full stdin pipe times out first. process.stdout = os.fdopen(stdout_read_fd, "rb", buffering=0) process.stderr = os.fdopen(stderr_read_fd, "rb", buffering=0) diff --git a/tests/test_shared_lean_runtime.py b/tests/test_shared_lean_runtime.py index 8b61bec5..b1b10774 100644 --- a/tests/test_shared_lean_runtime.py +++ b/tests/test_shared_lean_runtime.py @@ -3,11 +3,15 @@ from __future__ import annotations import asyncio +import io +import json +import os import socket import subprocess import sys import threading import time +from types import SimpleNamespace import pytest @@ -16,6 +20,9 @@ PROTOCOL_VERSION, LeanRuntimeClient, LeanRuntimeError, + LeanRuntimeOutcomeUnknown, + LeanRuntimeProtocolError, + LeanRuntimeRemoteError, LeanRuntimeUnavailable, ) from servers.lean_runtime import ( @@ -68,6 +75,9 @@ def run(self, code, **kwargs): def get_memory_usage(self): return 0.25 + def is_usable(self): + return not self._shutdown + def shutdown(self): self._shutdown = True @@ -118,10 +128,12 @@ def create_pool(root): assert first == second == "Compiles successfully" assert len(pools) == 1 - assert pools[0].calls == [ - ("#check Nat", {"timeout": 30.0}), - ("#check Int", {"timeout": 3.0}), - ] + assert [call[0] for call in pools[0].calls] == ["#check Nat", "#check Int"] + now = time.monotonic() + default_deadline = pools[0].calls[0][1]["deadline"] + explicit_deadline = pools[0].calls[1][1]["deadline"] + assert 0 < default_deadline - now <= 30.0 + assert 0 < explicit_deadline - now <= 3.0 warm = services.dispatch("repl.status", {"project_dir": str(project)}) assert warm["state"] == "warm" assert warm["memory_usage_gb"] == 0.25 @@ -131,6 +143,119 @@ def create_pool(root): assert pools[0]._shutdown is True +def test_repl_release_failure_after_result_is_outcome_unknown(tmp_path): + project = make_lake_project(tmp_path, "release-failure") + pools = [] + + class InvalidAfterRunPool(FakePool): + def is_usable(self): + if self.calls: + raise RuntimeError("post-result validation failed") + return True + + def create_pool(root): + pool = InvalidAfterRunPool(root) + pools.append(pool) + return pool + + services = LeanRuntimeServices( + runtime_config(), + repl_factory=create_pool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + try: + response = services.dispatch( + "repl.run", + {"project_dir": str(project), "code": "#check Nat", "timeout": 3}, + ) + finally: + services.close() + + assert pools[0].calls + assert "execution outcome unknown" in response + assert "must not be replayed" in response + assert "post-result validation failed" in response + + +def test_unexpected_pool_failure_after_attempt_is_outcome_unknown(tmp_path): + project = make_lake_project(tmp_path, "pool-failure") + calls = [] + + class FailingPool(FakePool): + def run(self, code, **kwargs): + calls.append((code, kwargs)) + raise RuntimeError("unexpected pool failure") + + services = LeanRuntimeServices( + runtime_config(), + repl_factory=FailingPool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + try: + response = services.dispatch( + "repl.run", + {"project_dir": str(project), "code": "#check Nat", "timeout": 3}, + ) + finally: + services.close() + + assert calls + assert "execution outcome unknown" in response + assert "must not be replayed" in response + assert "unexpected pool failure" in response + + +def test_pool_admission_failure_remains_retryable(tmp_path): + from servers.repl.pool import ReplPoolBusyError + + project = make_lake_project(tmp_path, "pool-busy") + + class BusyPool(FakePool): + def run(self, code, **kwargs): + raise ReplPoolBusyError("worker queue expired before dispatch") + + services = LeanRuntimeServices( + runtime_config(), + repl_factory=BusyPool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + try: + with pytest.raises(ReplPoolBusyError, match="before dispatch"): + services.dispatch( + "repl.run", + {"project_dir": str(project), "code": "#check Nat", "timeout": 3}, + ) + finally: + services.close() + + +def test_status_reports_an_active_poisoned_pool_as_retiring(tmp_path): + project = make_lake_project(tmp_path, "retiring-status") + services = LeanRuntimeServices( + runtime_config(), + repl_factory=FakePool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + try: + with services.repl_projects.lease(str(project)) as pool: + assert pool is not None + pool._shutdown = True + + status = services.dispatch( + "repl.status", + {"project_dir": str(project)}, + ) + + assert status["state"] == "retiring" + assert status["shutdown"] is True + finally: + services.close() + + def test_shared_runtime_disables_ambiguous_repl_retries(tmp_path, monkeypatch): from servers import lean_runtime @@ -281,6 +406,204 @@ def test_idle_ttl_never_closes_an_active_resource(tmp_path): cache.close() +def test_failed_retirement_blocks_replacement_without_losing_ownership(tmp_path): + first = make_lake_project(tmp_path, "retiring-first") + second = make_lake_project(tmp_path, "retiring-second") + created = [] + allow_close = False + + def factory(root): + created.append(root) + return root + + def close(resource): + if not allow_close: + raise RuntimeError("cleanup failed") + + cache = ProjectResourceCache( + factory, + close, + max_entries=1, + idle_seconds=1800, + start_sweeper=False, + ) + with cache.lease(str(first)): + pass + + with pytest.raises(ProjectResourceBusyError, match="failed to retire"): + with cache.lease(str(second)): + pytest.fail("replacement must wait for confirmed cleanup") + + assert created == [first.resolve()] + assert cache.state(str(first)) == "retiring" + assert cache.stats()["retiring"] == [str(first.resolve())] + + allow_close = True + with cache.lease(str(second)) as resource: + assert resource == second.resolve() + assert created == [first.resolve(), second.resolve()] + cache.close() + + +def test_concurrent_replacement_has_only_one_retirement_owner(tmp_path): + first = make_lake_project(tmp_path, "single-closer-first") + second = make_lake_project(tmp_path, "single-closer-second") + close_started = threading.Event() + release_close = threading.Event() + concurrent_close = threading.Event() + close_active = 0 + first_close_calls = 0 + + def close(resource): + nonlocal close_active, first_close_calls + if resource != first.resolve(): + return + first_close_calls += 1 + close_active += 1 + if close_active > 1: + concurrent_close.set() + close_started.set() + release_close.wait(timeout=2) + close_active -= 1 + + cache = ProjectResourceCache( + lambda root: root, + close, + max_entries=1, + idle_seconds=1800, + start_sweeper=False, + ) + with cache.lease(str(first)): + pass + + errors = [] + + def replace(): + try: + with cache.lease(str(second)): + pass + except BaseException as error: + errors.append(error) + + callers = [threading.Thread(target=replace) for _ in range(2)] + callers[0].start() + assert close_started.wait(timeout=1) + callers[1].start() + assert not concurrent_close.wait(timeout=0.1) + release_close.set() + for caller in callers: + caller.join(timeout=2) + + assert all(not caller.is_alive() for caller in callers) + assert errors == [] + assert first_close_calls == 1 + cache.close() + + +def test_cache_close_retains_failed_resources_for_a_later_retry(tmp_path): + project = make_lake_project(tmp_path, "close-retry") + close_calls = 0 + + def close(resource): + nonlocal close_calls + close_calls += 1 + if close_calls == 1: + raise RuntimeError("cleanup failed") + + cache = ProjectResourceCache( + lambda root: root, + close, + max_entries=1, + idle_seconds=1800, + start_sweeper=False, + ) + with cache.lease(str(project)): + pass + + with pytest.raises(RuntimeError, match="failed to retire 1"): + cache.close() + + assert cache.stats()["retiring"] == [str(project.resolve())] + cache.close() + assert close_calls == 2 + assert cache.stats()["retiring"] == [] + + +def test_lease_preserves_operation_cancellation_when_release_also_fails(tmp_path): + project = make_lake_project(tmp_path, "release-cancellation") + + def is_valid(resource): + raise asyncio.CancelledError("release") + + cache = ProjectResourceCache( + lambda root: root, + lambda resource: None, + max_entries=1, + idle_seconds=1800, + is_valid=is_valid, + start_sweeper=False, + ) + + with pytest.raises(KeyboardInterrupt, match="operation") as raised: + with cache.lease(str(project)): + raise KeyboardInterrupt("operation") + + if hasattr(raised.value, "add_note"): + assert raised.value.__notes__ == [ + "Lean project resource release also failed: release" + ] + cache.close() + + +def test_services_attempt_lsp_cleanup_after_repl_cleanup_failure(monkeypatch): + services = LeanRuntimeServices( + runtime_config(), + repl_factory=FakePool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + lsp_closed = [] + monkeypatch.setattr( + services.repl_projects, + "close", + lambda: (_ for _ in ()).throw(RuntimeError("REPL cleanup failed")), + ) + monkeypatch.setattr( + services.lsp_projects, + "close", + lambda: lsp_closed.append(True), + ) + + with pytest.raises(RuntimeError, match="REPL cleanup failed"): + services.close() + + assert lsp_closed == [True] + + +def test_terminal_cleanup_retries_without_releasing_ownership(monkeypatch): + from servers import lean_runtime + + close_calls = 0 + delays = [] + + class Services: + def close(self): + nonlocal close_calls + close_calls += 1 + if close_calls < 3: + raise RuntimeError("cleanup failed") + + monkeypatch.setattr(lean_runtime.time, "sleep", delays.append) + + lean_runtime._close_services_until_clean(Services()) + + assert close_calls == 3 + assert delays == [ + lean_runtime.TERMINAL_CLEANUP_RETRY_SECONDS, + lean_runtime.TERMINAL_CLEANUP_RETRY_SECONDS * 2, + ] + + def test_stdio_mcp_adapters_delegate_without_owning_lean_state(): from servers.lsp.server import create_lsp_server from servers.repl.server import create_repl_server @@ -346,98 +669,892 @@ def request(self, method, params): ] -def test_lsp_diagnostic_formatting_remains_stable(): - from servers.lsp.server import format_lsp_diagnostics - - assert format_lsp_diagnostics([]).startswith("No diagnostics") - formatted = format_lsp_diagnostics( - [ - { - "severity": 1, - "message": "unknown identifier", - "range": {"start": {"line": 2, "character": 4}}, - } - ] - ) - assert formatted == ( - "Diagnostics: 1 error(s), 0 warning(s)\n" - "3:4: error: unknown identifier" - ) +@pytest.mark.parametrize( + ("timeout", "configured_default", "expected_deadline"), + [(5, None, 105.0), (None, "7", 107.0)], +) +def test_repl_client_reuses_one_deadline_across_autostart_retry( + runtime_dir, + monkeypatch, + timeout, + configured_default, + expected_deadline, +): + from servers import lean_client + now = [100.0] + attempts = [] + startup_deadlines = [] + client = LeanRuntimeClient(socket_path=runtime_dir / "deadline.sock") + if configured_default is not None: + monkeypatch.setenv("AUTOFORM_REPL_REQUEST_TIMEOUT", configured_default) + + def request_once(method, params, **kwargs): + attempts.append(kwargs["deadline"]) + if len(attempts) == 1: + raise LeanRuntimeUnavailable("not listening") + return "Compiles successfully" + + def ensure_running(*, deadline): + startup_deadlines.append(deadline) + now[0] = 102.0 + return {"running": True} + + monkeypatch.setattr(lean_client.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(client, "_request_once", request_once) + monkeypatch.setattr(client, "ensure_running", ensure_running) + + assert client.request( + "repl.run", + {"project_dir": "/lean", "code": "#check Nat", "timeout": timeout}, + ) == "Compiles successfully" + assert attempts == [expected_deadline, expected_deadline] + assert startup_deadlines == [expected_deadline] + + +def test_repl_wire_deadline_is_internal_and_response_allows_cleanup( + runtime_dir, + monkeypatch, +): + from servers import lean_client -def test_concurrent_clients_boot_one_daemon_that_outlives_each_client(runtime_dir, monkeypatch): - socket_path = runtime_dir / "lean.sock" - monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "1") - monkeypatch.setenv("AUTOFORM_MAX_LEAN_PROJECTS", "1") - clients = [ - LeanRuntimeClient(socket_path=socket_path, startup_timeout=15), - LeanRuntimeClient(socket_path=socket_path, startup_timeout=15), - ] - barrier = threading.Barrier(3) - pids = [] - errors = [] + sent = [] + socket_timeouts = [] - def start(client): - barrier.wait() - try: - pids.append(client.ensure_running()["pid"]) - except BaseException as error: - errors.append(error) + class RespondingSocket: + def settimeout(self, timeout): + socket_timeouts.append(timeout) - threads = [threading.Thread(target=start, args=(client,)) for client in clients] - for thread in threads: - thread.start() - barrier.wait() - for thread in threads: - thread.join(timeout=20) + def connect(self, path): + pass - try: - assert errors == [] - assert all(not thread.is_alive() for thread in threads) - assert len(pids) == 2 - assert len(set(pids)) == 1 + def sendall(self, payload): + sent.append(json.loads(payload)) + + def recv(self, size): + return json.dumps( + { + "v": PROTOCOL_VERSION, + "id": sent[0]["id"], + "ok": True, + "result": "Compiles successfully", + } + ).encode() + b"\n" - # Clients own no process handle or shutdown hook. Losing the client that - # happened to bootstrap the daemon cannot stop shared Lean state. - del clients[0] - assert clients[0].ping()["pid"] == pids[0] - finally: - try: - clients[-1].stop() - except LeanRuntimeUnavailable: + def close(self): pass - deadline = time.monotonic() + 5 - while socket_path.exists() and time.monotonic() < deadline: - time.sleep(0.025) - assert not socket_path.exists() + monkeypatch.setattr(lean_client.time, "monotonic", lambda: 100.0) + monkeypatch.setattr(lean_client.socket, "socket", lambda *args: RespondingSocket()) + client = LeanRuntimeClient( + socket_path=runtime_dir / "wire-deadline.sock", + response_timeout=900, + ) + params = {"project_dir": "/lean", "code": "#check Nat", "timeout": 5} + + assert client.request("repl.run", params, autostart=False) == "Compiles successfully" + assert sent[0]["deadline"] == 105.0 + assert sent[0]["params"] == params + assert "deadline" not in sent[0]["params"] + assert socket_timeouts == [ + 2.0, + 5.0, + 5.0 + lean_client.REPL_RESPONSE_GRACE_SECONDS, + ] -def test_daemon_outlives_the_separate_process_that_started_it( - tmp_path, +def test_repl_response_budget_must_cover_operation_and_cleanup_before_dispatch( runtime_dir, - repo_root, monkeypatch, ): - socket_path = runtime_dir / "owner.sock" - project = make_lake_project(tmp_path, "cold") - monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "1") - monkeypatch.setenv("AUTOFORM_MAX_LEAN_PROJECTS", "1") - helper = subprocess.run( - [ - sys.executable, - "-c", - ( - "import sys; " - "from servers.lean_client import LeanRuntimeClient; " - "print(LeanRuntimeClient(socket_path=sys.argv[1]).ensure_running()['pid'])" - ), - str(socket_path), - ], - cwd=repo_root, - capture_output=True, - text=True, - timeout=20, + from servers import lean_client + + monkeypatch.setattr( + lean_client.socket, + "socket", + lambda *args: pytest.fail("an invalid response budget must not connect"), + ) + client = LeanRuntimeClient( + socket_path=runtime_dir / "short-response.sock", + response_timeout=36, + ) + + with pytest.raises(LeanRuntimeError, match="response timeout must exceed"): + client.request( + "repl.run", + {"project_dir": "/lean", "code": "#check Nat", "timeout": 5}, + ) + + +@pytest.mark.parametrize("deadline", [True, "soon", float("nan"), float("inf")]) +def test_invalid_repl_client_deadline_never_connects( + runtime_dir, + monkeypatch, + deadline, +): + from servers import lean_client + + monkeypatch.setattr( + lean_client.socket, + "socket", + lambda *args: pytest.fail("an invalid deadline must not connect"), + ) + client = LeanRuntimeClient(socket_path=runtime_dir / "invalid-deadline.sock") + + with pytest.raises(LeanRuntimeError, match="deadline must be a finite number"): + client.request( + "repl.run", + {"project_dir": "/lean", "code": "#check Nat", "timeout": 5}, + deadline=deadline, + ) + + +def test_repl_response_read_rechecks_one_absolute_deadline( + runtime_dir, + monkeypatch, +): + from servers import lean_client + + now = [100.0] + receives = [] + + class DribblingSocket: + def settimeout(self, timeout): + pass + + def connect(self, path): + pass + + def sendall(self, payload): + pass + + def recv(self, size): + receives.append(size) + now[0] = 134.0 + return b"{" + + def close(self): + pass + + monkeypatch.setattr(lean_client.time, "monotonic", lambda: now[0]) + monkeypatch.setattr( + lean_client.socket, + "socket", + lambda *args: DribblingSocket(), + ) + client = LeanRuntimeClient(socket_path=runtime_dir / "dribble.sock") + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed"): + client.request( + "repl.run", + {"project_dir": "/lean", "code": "#check Nat", "timeout": 1}, + autostart=False, + ) + + assert len(receives) == 1 + + +def test_repl_response_arriving_after_deadline_is_not_accepted( + runtime_dir, + monkeypatch, +): + from servers import lean_client + + now = [100.0] + sent = [] + + class LateResponseSocket: + def settimeout(self, timeout): + pass + + def connect(self, path): + pass + + def sendall(self, payload): + sent.append(json.loads(payload)) + + def recv(self, size): + now[0] = 134.0 + return json.dumps( + { + "v": PROTOCOL_VERSION, + "id": sent[0]["id"], + "ok": True, + "result": "late result", + } + ).encode() + b"\n" + + def close(self): + pass + + monkeypatch.setattr(lean_client.time, "monotonic", lambda: now[0]) + monkeypatch.setattr( + lean_client.socket, + "socket", + lambda *args: LateResponseSocket(), + ) + client = LeanRuntimeClient(socket_path=runtime_dir / "late-response.sock") + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed"): + client.request( + "repl.run", + {"project_dir": "/lean", "code": "#check Nat", "timeout": 1}, + autostart=False, + ) + + +def test_expired_repl_deadline_is_not_dispatched_after_autostart( + runtime_dir, + monkeypatch, +): + from servers import lean_client + + now = [100.0] + connects = [] + sends = [] + + class MissingThenForbiddenSocket: + def __init__(self, number): + self.number = number + + def settimeout(self, timeout): + pass + + def connect(self, path): + connects.append(self.number) + if self.number == 1: + raise FileNotFoundError(path) + pytest.fail("an expired request must not reconnect") + + def sendall(self, payload): + sends.append(payload) + + def close(self): + pass + + sockets = [] + + def create_socket(*args): + candidate = MissingThenForbiddenSocket(len(sockets) + 1) + sockets.append(candidate) + return candidate + + def ensure_running(*, deadline): + assert deadline == 105.0 + now[0] = 106.0 + return {"running": True} + + monkeypatch.setattr(lean_client.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(lean_client.socket, "socket", create_socket) + client = LeanRuntimeClient(socket_path=runtime_dir / "expired.sock") + monkeypatch.setattr(client, "ensure_running", ensure_running) + + with pytest.raises(LeanRuntimeUnavailable, match="expired before dispatch"): + client.request( + "repl.run", + {"project_dir": "/lean", "code": "#check Nat", "timeout": 5}, + ) + + assert connects == [1] + assert sends == [] + + +@pytest.mark.parametrize( + ("client_deadline", "server_timeout", "expected_deadline"), + [ + (110.0, 30.0, 110.0), + (200.0, 5.0, 105.0), + ], +) +def test_runtime_caps_client_deadline_and_spends_admission_time( + tmp_path, + monkeypatch, + client_deadline, + server_timeout, + expected_deadline, +): + from servers import lean_runtime + + project = make_lake_project(tmp_path, "runtime-deadline") + now = [100.0] + pool = FakePool(project.resolve()) + lease_deadlines = [] + + class Lease: + def __enter__(self): + now[0] = 103.0 + return pool + + def __exit__(self, *args): + return False + + class Projects: + def lease(self, project_dir, **kwargs): + lease_deadlines.append(kwargs["deadline"]) + return Lease() + + monkeypatch.setattr(lean_runtime.time, "monotonic", lambda: now[0]) + services = LeanRuntimeServices( + runtime_config(), + repl_factory=FakePool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + original_projects = services.repl_projects + services.repl_projects = Projects() + try: + assert services.dispatch( + "repl.run", + { + "project_dir": str(project), + "code": "#check Nat", + "timeout": server_timeout, + }, + client_deadline=client_deadline, + ) == "Compiles successfully" + finally: + services.repl_projects = original_projects + services.close() + + assert lease_deadlines == [expected_deadline] + assert pool.calls == [ + ("#check Nat", {"deadline": expected_deadline}) + ] + + +def test_expired_runtime_deadline_never_warms_or_dispatches_a_pool(tmp_path): + project = make_lake_project(tmp_path, "expired-runtime") + pools = [] + services = LeanRuntimeServices( + runtime_config(), + repl_factory=lambda root: pools.append(FakePool(root)) or pools[-1], + lsp_factory=FakeLsp, + start_sweepers=False, + ) + try: + with pytest.raises(ProjectResourceBusyError, match="expired before admission"): + services.dispatch( + "repl.run", + {"project_dir": str(project), "code": "#check Nat", "timeout": 30}, + client_deadline=time.monotonic() - 1, + ) + assert pools == [] + finally: + services.close() + + +def test_runtime_does_not_dispatch_when_deadline_expires_during_admission( + tmp_path, + monkeypatch, +): + from servers import lean_runtime + + project = make_lake_project(tmp_path, "expired-admission") + now = [100.0] + pool = FakePool(project.resolve()) + + class Lease: + def __enter__(self): + now[0] = 106.0 + return pool + + def __exit__(self, *args): + return False + + class Projects: + def lease(self, project_dir, **kwargs): + assert kwargs["deadline"] == 105.0 + return Lease() + + monkeypatch.setattr(lean_runtime.time, "monotonic", lambda: now[0]) + services = LeanRuntimeServices( + runtime_config(), + repl_factory=FakePool, + lsp_factory=FakeLsp, + start_sweepers=False, + ) + original_projects = services.repl_projects + services.repl_projects = Projects() + try: + with pytest.raises(ProjectResourceBusyError, match="expired before execution"): + services.dispatch( + "repl.run", + {"project_dir": str(project), "code": "#check Nat", "timeout": 30}, + client_deadline=105.0, + ) + finally: + services.repl_projects = original_projects + services.close() + + assert pool.calls == [] + + +def test_cache_does_not_start_a_resource_after_its_absolute_deadline(tmp_path): + project = make_lake_project(tmp_path, "expired-cache-start") + now = iter((0.0, 2.0)) + created = [] + cache = ProjectResourceCache( + lambda root: created.append(root) or root, + lambda resource: None, + max_entries=1, + idle_seconds=1800, + start_sweeper=False, + clock=lambda: next(now), + ) + + with pytest.raises(ProjectResourceBusyError, match="not enough response budget"): + with cache.lease(str(project), deadline=1.0): + pytest.fail("an expired lease must not start a project resource") + + assert created == [] + cache.close() + + +class _RuntimeRequestSocket: + def __init__(self, payload): + self.payload = payload + self.timeouts = [] + + def settimeout(self, timeout): + self.timeouts.append(timeout) + + def recv(self, size): + payload, self.payload = self.payload, b"" + return payload + + +@pytest.mark.parametrize("include_deadline", [False, True]) +def test_runtime_wire_deadline_is_optional_and_outside_params(include_deadline): + from servers.lean_runtime import LeanRuntimeRequestHandler + + calls = [] + + class Services: + def dispatch(self, method, params, *, client_deadline=None): + calls.append((method, params, client_deadline)) + return "Compiles successfully" + + request = { + "v": PROTOCOL_VERSION, + "id": "request-id", + "method": "repl.run", + "params": {"project_dir": "/lean", "code": "#check Nat", "timeout": 5}, + } + deadline = time.monotonic() + 60 if include_deadline else None + if include_deadline: + request["deadline"] = deadline + handler = object.__new__(LeanRuntimeRequestHandler) + request_socket = _RuntimeRequestSocket(json.dumps(request).encode() + b"\n") + handler.request = request_socket + handler.wfile = io.BytesIO() + services = Services() + services.config = SimpleNamespace(rpc_read_timeout=1.0) + handler.server = SimpleNamespace( + services=services, + request_shutdown=lambda: None, + ) + + handler.handle() + + assert calls == [("repl.run", request["params"], deadline)] + response = json.loads(handler.wfile.getvalue()) + assert response["ok"] is True + assert request_socket.timeouts[-1] > request_socket.timeouts[0] + + +@pytest.mark.parametrize( + "raw", + [ + b'{"v":1,"v":1,"id":"request-id","method":"daemon.ping","params":{}}\n', + b'{"v":true,"id":"request-id","method":"daemon.ping","params":{}}\n', + b'{"v":1,"id":1,"method":"daemon.ping","params":{}}\n', + b'{"v":1,"id":"","method":"daemon.ping","params":{}}\n', + b'{"v":1,"id":"request-id","method":"daemon.ping","params":{},"extra":1}\n', + b'{"v":1,"id":"request-id","method":"daemon.ping","params":{},"deadline":1}\n', + b'{"v":1,"id":"request-id","method":"repl.run","params":{},"deadline":NaN}\n', + ], +) +def test_runtime_server_rejects_malformed_request_envelopes(raw): + from servers.lean_runtime import LeanRuntimeRequestHandler + + class Services: + config = SimpleNamespace(rpc_read_timeout=1.0) + + def dispatch(self, *args, **kwargs): + pytest.fail("a malformed request must not be dispatched") + + handler = object.__new__(LeanRuntimeRequestHandler) + handler.request = _RuntimeRequestSocket(raw) + handler.wfile = io.BytesIO() + handler.server = SimpleNamespace( + services=Services(), + request_shutdown=lambda: None, + ) + + handler.handle() + + response = json.loads(handler.wfile.getvalue()) + assert response["ok"] is False + assert response["error"]["type"] == "ValueError" + + +def test_oversized_repl_result_preserves_unknown_outcome( + monkeypatch, +): + from servers import lean_runtime + + class Services: + config = SimpleNamespace(rpc_read_timeout=1.0) + + def dispatch(self, method, params, *, client_deadline=None): + return "x" * 1_000 + + request = { + "v": PROTOCOL_VERSION, + "id": "request-id", + "method": "repl.run", + "params": {"project_dir": "/lean", "code": "#check Nat"}, + } + handler = object.__new__(lean_runtime.LeanRuntimeRequestHandler) + handler.request = _RuntimeRequestSocket(json.dumps(request).encode() + b"\n") + handler.wfile = io.BytesIO() + handler.server = SimpleNamespace( + services=Services(), + request_shutdown=lambda: None, + ) + monkeypatch.setattr(lean_runtime, "MAX_MESSAGE_BYTES", 512) + + handler.handle() + + response = json.loads(handler.wfile.getvalue()) + assert response["ok"] is True + assert "execution outcome unknown" in response["result"] + assert "must not be replayed" in response["result"] + + +def test_runtime_server_request_read_rechecks_one_absolute_deadline(monkeypatch): + from servers import lean_runtime + + now = [100.0] + receives = [] + + class DribblingRequest: + def settimeout(self, timeout): + pass + + def recv(self, size): + receives.append(size) + now[0] = 101.0 + return b"{" + + handler = object.__new__(lean_runtime.LeanRuntimeRequestHandler) + handler.request = DribblingRequest() + handler.server = SimpleNamespace( + services=SimpleNamespace( + config=SimpleNamespace(rpc_read_timeout=0.5) + ) + ) + monkeypatch.setattr(lean_runtime.time, "monotonic", lambda: now[0]) + + with pytest.raises(TimeoutError, match="read deadline expired"): + handler._read_request() + + assert len(receives) == 1 + + +def test_runtime_server_resets_write_timeout_after_a_dribbling_request( + monkeypatch, + caplog, +): + from servers import lean_runtime + + now = [100.0] + request = json.dumps( + { + "v": PROTOCOL_VERSION, + "id": "request-id", + "method": "daemon.ping", + "params": {}, + } + ).encode() + b"\n" + + class DribblingRequest: + def __init__(self): + self.chunks = [request[:1], request[1:]] + self.timeouts = [] + + def settimeout(self, timeout): + self.timeouts.append(timeout) + + def recv(self, size): + chunk = self.chunks.pop(0) + if self.chunks: + now[0] = 100.99 + return chunk + + class TimingOutWriter: + def __init__(self, request_socket): + self.request_socket = request_socket + self.response = None + self.write_timeout = None + + def write(self, encoded): + self.response = json.loads(encoded) + self.write_timeout = self.request_socket.timeouts[-1] + raise socket.timeout("blocked") + + def flush(self): + pytest.fail("flush must not run after a failed write") + + calls = [] + + class Services: + config = SimpleNamespace(rpc_read_timeout=1.0, response_timeout=900.0) + + def dispatch(self, method, params, *, client_deadline=None): + calls.append((method, params, client_deadline)) + return {"pid": 1} + + request_socket = DribblingRequest() + writer = TimingOutWriter(request_socket) + handler = object.__new__(lean_runtime.LeanRuntimeRequestHandler) + handler.request = request_socket + handler.wfile = writer + handler.server = SimpleNamespace( + services=Services(), + request_shutdown=lambda: None, + ) + monkeypatch.setattr(lean_runtime.time, "monotonic", lambda: now[0]) + + handler.handle() + + assert calls == [("daemon.ping", {}, None)] + assert request_socket.timeouts[-2] == pytest.approx(0.01) + assert writer.write_timeout == pytest.approx(lean_runtime.RUNTIME_SAFETY_SECONDS) + assert writer.response == { + "v": PROTOCOL_VERSION, + "id": "request-id", + "ok": True, + "result": {"pid": 1}, + } + assert "disconnected before receiving" in caplog.text + + +def test_lsp_diagnostic_formatting_remains_stable(): + from servers.lsp.server import format_lsp_diagnostics + + assert format_lsp_diagnostics([]).startswith("No diagnostics") + formatted = format_lsp_diagnostics( + [ + { + "severity": 1, + "message": "unknown identifier", + "range": {"start": {"line": 2, "character": 4}}, + } + ] + ) + assert formatted == ( + "Diagnostics: 1 error(s), 0 warning(s)\n" + "3:4: error: unknown identifier" + ) + + +def test_startup_times_out_while_previous_runtime_retains_lifetime_lock( + runtime_dir, + monkeypatch, +): + import fcntl + + socket_path = runtime_dir / "retiring.sock" + client = LeanRuntimeClient(socket_path=socket_path, startup_timeout=0.01) + lock_calls = 0 + + def flock(fd, operation): + nonlocal lock_calls + lock_calls += 1 + if lock_calls == 1: + return + raise BlockingIOError + + monkeypatch.setattr(fcntl, "flock", flock) + + with pytest.raises(LeanRuntimeUnavailable, match="still be cleaning up"): + client.ensure_running() + + +def test_startup_does_not_acquire_a_free_lock_after_its_deadline( + runtime_dir, + monkeypatch, +): + from servers import lean_client + + client = LeanRuntimeClient( + socket_path=runtime_dir / "expired.sock", + startup_timeout=1, + ) + now = [100.0] + lock_calls = [] + + def unavailable_ping(*, autostart=False, deadline=None): + now[0] = 102.0 + raise LeanRuntimeUnavailable("not listening") + + monkeypatch.setattr(lean_client.time, "monotonic", lambda: now[0]) + monkeypatch.setattr(client, "ping", unavailable_ping) + monkeypatch.setattr( + "fcntl.flock", + lambda *args: lock_calls.append(args), + ) + + with pytest.raises(LeanRuntimeUnavailable, match="startup coordination"): + client.ensure_running() + + assert lock_calls == [] + + +def test_previous_build_shutdown_uses_the_startup_deadline(runtime_dir, monkeypatch): + from servers import lean_client + + current_socket = runtime_dir / f"lean-v{PROTOCOL_VERSION}-{INSTALL_PATH_ID}-new.sock" + old_socket = runtime_dir / f"lean-v{PROTOCOL_VERSION}-{INSTALL_PATH_ID}-old.sock" + old_socket.touch() + client = LeanRuntimeClient(socket_path=current_socket) + client._uses_default_paths = True + calls = [] + + class PreviousClient: + def __init__(self, *, socket_path, **kwargs): + assert socket_path == old_socket + + def request(self, method, *, autostart, deadline): + calls.append((method, deadline)) + return {"build_generation": 0} + + def stop(self, *, deadline): + calls.append(("stop", deadline)) + return {"pid": 7} + + monkeypatch.setattr(lean_client, "LeanRuntimeClient", PreviousClient) + + assert client._stop_previous_builds(deadline=123.0) == [7] + assert calls == [("daemon.ping", 123.0), ("stop", 123.0)] + + +def test_build_generation_includes_shared_environment_code(monkeypatch): + from servers import lean_client + + environment_module = lean_client.PACKAGE_ROOT / "servers" / "__init__.py" + + def fake_stat(path): + return SimpleNamespace( + st_mtime_ns=2 if path == environment_module else 1 + ) + + monkeypatch.setattr(lean_client.Path, "stat", fake_stat) + + assert lean_client._build_generation() == 2 + + +def test_failed_start_never_hard_kills_a_daemon_that_may_own_work(runtime_dir): + client = LeanRuntimeClient(socket_path=runtime_dir / "failed-start.sock") + + class Process: + returncode = None + terminate_calls = 0 + kill_calls = 0 + + def poll(self): + return None + + def terminate(self): + self.terminate_calls += 1 + + def wait(self, timeout): + raise subprocess.TimeoutExpired("runtime", timeout) + + def kill(self): + self.kill_calls += 1 + + process = Process() + client._terminate_failed_start(process) + + assert process.terminate_calls == 1 + assert process.kill_calls == 0 + + +def test_concurrent_clients_boot_one_daemon_that_outlives_each_client(runtime_dir, monkeypatch): + socket_path = runtime_dir / "lean.sock" + monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "1") + monkeypatch.setenv("AUTOFORM_MAX_LEAN_PROJECTS", "1") + clients = [ + LeanRuntimeClient(socket_path=socket_path, startup_timeout=15), + LeanRuntimeClient(socket_path=socket_path, startup_timeout=15), + ] + barrier = threading.Barrier(3) + pids = [] + errors = [] + + def start(client): + barrier.wait() + try: + pids.append(client.ensure_running()["pid"]) + except BaseException as error: + errors.append(error) + + threads = [threading.Thread(target=start, args=(client,)) for client in clients] + for thread in threads: + thread.start() + barrier.wait() + for thread in threads: + thread.join(timeout=20) + + try: + assert errors == [] + assert all(not thread.is_alive() for thread in threads) + assert len(pids) == 2 + assert len(set(pids)) == 1 + + # Clients own no process handle or shutdown hook. Losing the client that + # happened to bootstrap the daemon cannot stop shared Lean state. + del clients[0] + assert clients[0].ping()["pid"] == pids[0] + finally: + try: + clients[-1].stop() + except LeanRuntimeUnavailable: + pass + + deadline = time.monotonic() + 5 + while socket_path.exists() and time.monotonic() < deadline: + time.sleep(0.025) + assert not socket_path.exists() + + +def test_daemon_outlives_the_separate_process_that_started_it( + tmp_path, + runtime_dir, + repo_root, + monkeypatch, +): + socket_path = runtime_dir / "owner.sock" + project = make_lake_project(tmp_path, "cold") + monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "1") + monkeypatch.setenv("AUTOFORM_MAX_LEAN_PROJECTS", "1") + helper = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; " + "from servers.lean_client import LeanRuntimeClient; " + "print(LeanRuntimeClient(socket_path=sys.argv[1]).ensure_running()['pid'])" + ), + str(socket_path), + ], + cwd=repo_root, + capture_output=True, + text=True, + timeout=20, ) assert helper.returncode == 0, helper.stderr owner_pid = int(helper.stdout.strip()) @@ -477,6 +1594,8 @@ def test_stop_then_immediate_start_is_serialized(runtime_dir, monkeypatch): def test_new_build_replaces_previous_runtime_at_same_install_path(runtime_dir, monkeypatch): + import fcntl + monkeypatch.setenv("AUTOFORM_RUNTIME_DIR", str(runtime_dir)) monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "1") old_socket = runtime_dir / f"lean-v{PROTOCOL_VERSION}-{INSTALL_PATH_ID}-old.sock" @@ -488,6 +1607,12 @@ def test_new_build_replaces_previous_runtime_at_same_install_path(runtime_dir, m current_pid = current.ensure_running()["pid"] assert current_pid != old_pid assert not old_socket.exists() + lifetime_fd = os.open(current.paths.lifetime_lock, os.O_RDWR) + try: + with pytest.raises(BlockingIOError): + fcntl.flock(lifetime_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + finally: + os.close(lifetime_fd) finally: current.stop() @@ -562,10 +1687,270 @@ def close(self): lambda: pytest.fail("an ambiguously dispatched request must not be retried"), ) - with pytest.raises(LeanRuntimeError, match="after request dispatch"): + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed"): + client.request("repl.run", {"project_dir": "/lean", "code": "#check Nat"}) + + +def test_post_dispatch_timeout_is_explicitly_outcome_unknown(runtime_dir, monkeypatch): + from servers import lean_client + + class TimingOutSocket: + def settimeout(self, timeout): + pass + + def connect(self, path): + pass + + def sendall(self, payload): + pass + + def recv(self, size): + raise socket.timeout + + def close(self): + pass + + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + monkeypatch.setattr(lean_client.socket, "socket", lambda *args: TimingOutSocket()) + monkeypatch.setattr( + client, + "ensure_running", + lambda **kwargs: pytest.fail("a dispatched request must not be retried"), + ) + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed"): + client.request("repl.run", {"project_dir": "/lean", "code": "#check Nat"}) + + +@pytest.mark.parametrize("response", [b"", b"{\n"]) +def test_post_dispatch_invalid_response_is_explicitly_outcome_unknown( + runtime_dir, + monkeypatch, + response, +): + from servers import lean_client + + class InvalidResponseSocket: + def settimeout(self, timeout): + pass + + def connect(self, path): + pass + + def sendall(self, payload): + pass + + def recv(self, size): + return response + + def close(self): + pass + + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + monkeypatch.setattr( + lean_client.socket, + "socket", + lambda *args: InvalidResponseSocket(), + ) + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed"): client.request("repl.run", {"project_dir": "/lean", "code": "#check Nat"}) +class _RuntimeResponseSocket: + def __init__(self, response): + self.response = response + + def settimeout(self, timeout): + pass + + def connect(self, path): + pass + + def sendall(self, payload): + pass + + def recv(self, size): + response, self.response = self.response, b"" + return response + + def close(self): + pass + + +def test_runtime_response_limit_counts_the_frame_delimiter(monkeypatch): + from servers import lean_client + + monkeypatch.setattr(lean_client, "MAX_MESSAGE_BYTES", 8) + connection = _RuntimeResponseSocket(b"12345678\n") + + with pytest.raises(LeanRuntimeProtocolError, match="message limit"): + LeanRuntimeClient._read_line( + connection, + deadline=time.monotonic() + 1, + ) + + +def _runtime_response(client, monkeypatch, response, *, method="repl.run"): + from servers import lean_client + + monkeypatch.setattr(lean_client.uuid, "uuid4", lambda: type("UUID", (), {"hex": "request-id"})()) + monkeypatch.setattr( + lean_client.socket, + "socket", + lambda *args: _RuntimeResponseSocket(response), + ) + return client.request(method, {"project_dir": "/lean", "code": "#check Nat"}) + + +@pytest.mark.parametrize( + ("method", "encoded_result", "expected"), + [ + ("repl.run", b'"Compiles successfully"', "Compiles successfully"), + ("lsp.diagnostics", b'"No diagnostics"', "No diagnostics"), + ("lsp.hover", b'"Nat"', "Nat"), + ("daemon.ping", b'{"running":true}', {"running": True}), + ("daemon.status", b'{"running":true}', {"running": True}), + ("daemon.shutdown", b'{"stopping":true}', {"stopping": True}), + ("repl.status", b'{"state":"cold"}', {"state": "cold"}), + ], +) +def test_runtime_response_accepts_declared_method_result( + runtime_dir, + monkeypatch, + method, + encoded_result, + expected, +): + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + + response = ( + b'{"v":1,"id":"request-id","ok":true,"result":' + + encoded_result + + b"}\n" + ) + + assert _runtime_response(client, monkeypatch, response, method=method) == expected + + +@pytest.mark.parametrize( + ("method", "encoded_result"), + [ + ("repl.run", b"null"), + ("repl.run", b"false"), + ("repl.run", b"0"), + ("repl.run", b"[1,2]"), + ("repl.run", b'{"nested":true}'), + ("daemon.ping", b'"running"'), + ], +) +def test_runtime_response_rejects_wrong_method_result_type( + runtime_dir, + monkeypatch, + method, + encoded_result, +): + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + response = ( + b'{"v":1,"id":"request-id","ok":true,"result":' + + encoded_result + + b"}\n" + ) + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed") as caught: + _runtime_response(client, monkeypatch, response, method=method) + + assert isinstance(caught.value.__cause__, LeanRuntimeProtocolError) + + +def test_runtime_response_accepts_exact_error_envelope(runtime_dir, monkeypatch): + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + + with pytest.raises(LeanRuntimeRemoteError, match="ValueError: bad request"): + _runtime_response( + client, + monkeypatch, + b'{"v":1,"id":"request-id","ok":false,' + b'"error":{"type":"ValueError","message":"bad request"}}\n', + ) + + +@pytest.mark.parametrize( + "response", + [ + b'{"v":true,"id":"request-id","ok":true,"result":null}\n', + b'{"v":1.0,"id":"request-id","ok":true,"result":null}\n', + b'{"v":1,"id":1,"ok":true,"result":null}\n', + b'{"v":1,"id":"wrong-id","ok":true,"result":null}\n', + b'{"v":1,"id":"request-id","ok":1,"result":null}\n', + b'{"v":1,"id":"request-id","ok":true}\n', + b'{"v":1,"id":"request-id","ok":true,"result":null,"error":null}\n', + b'{"v":1,"id":"request-id","ok":true,"result":null,"extra":null}\n', + b'{"v":1,"id":"request-id","ok":false}\n', + b'{"v":1,"id":"request-id","ok":false,"error":{},"result":null}\n', + b'{"v":1,"id":"request-id","ok":false,"error":{"type":"ValueError"}}\n', + b'{"v":1,"id":"request-id","ok":false,' + b'"error":{"type":"ValueError","message":"bad","extra":null}}\n', + b'{"v":1,"id":"request-id","ok":false,' + b'"error":{"type":1,"message":"bad"}}\n', + b'{"v":1,"id":"request-id","ok":false,' + b'"error":{"type":"ValueError","message":null}}\n', + ], +) +def test_runtime_response_rejects_malformed_envelopes( + runtime_dir, + monkeypatch, + response, +): + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed") as caught: + _runtime_response(client, monkeypatch, response) + + assert isinstance(caught.value.__cause__, LeanRuntimeProtocolError) + + +@pytest.mark.parametrize( + "response", + [ + b'{"v":1,"v":1,"id":"request-id","ok":true,"result":null}\n', + b'{"v":1,"id":"request-id","ok":true,"result":NaN}\n', + b'{"v":1,"id":"request-id","ok":true,"result":1e100000}\n', + b'{"v":1,"id":"request-id","ok":true,"result":' + + b"9" * 4_301 + + b"}\n", + b'{"v":1,"id":"request-id","ok":true,"result":' + + b"[" * 2_000 + + b"0" + + b"]" * 2_000 + + b"}\n", + ], +) +def test_runtime_response_rejects_noncanonical_json( + runtime_dir, + monkeypatch, + response, +): + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed") as caught: + _runtime_response(client, monkeypatch, response) + + assert isinstance(caught.value.__cause__, LeanRuntimeProtocolError) + + +def test_runtime_response_rejects_invalid_utf8_as_unknown_outcome( + runtime_dir, + monkeypatch, +): + client = LeanRuntimeClient(socket_path=runtime_dir / "fake.sock") + + with pytest.raises(LeanRuntimeOutcomeUnknown, match="must not be replayed") as caught: + _runtime_response(client, monkeypatch, b"\xff\n") + + assert isinstance(caught.value.__cause__, LeanRuntimeProtocolError) + + @pytest.mark.parametrize( ("name", "value", "match"), [ @@ -589,11 +1974,32 @@ def test_per_project_workers_cannot_exceed_node_budget(monkeypatch): LeanRuntimeConfig.from_environment() -def test_response_budget_includes_replacement_and_failed_pool_cleanup(monkeypatch): +def test_response_budget_does_not_scale_with_cold_repl_pool_size(monkeypatch): monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "3") monkeypatch.setenv("AUTOFORM_REPL_WORKERS_PER_PROJECT", "3") monkeypatch.setenv("AUTOFORM_RUNTIME_RESPONSE_TIMEOUT", "860") - with pytest.raises(ValueError, match="REPL worker startup"): + + config = LeanRuntimeConfig.from_environment() + + assert config.repl_workers_per_project == 3 + assert config.response_timeout == 860 + + +def test_client_repl_response_grace_matches_daemon_cleanup_reserve(): + from servers import lean_client, lean_runtime + from servers.repl.pool import DEFAULT_POOL_CLEANUP_SECONDS + + assert lean_client.REPL_RESPONSE_GRACE_SECONDS == ( + DEFAULT_POOL_CLEANUP_SECONDS + lean_runtime.RUNTIME_SAFETY_SECONDS + ) + + +def test_response_budget_must_leave_room_for_repl_cleanup(monkeypatch): + monkeypatch.setenv("AUTOFORM_RUNTIME_RESPONSE_TIMEOUT", "272") + monkeypatch.setenv("LEAN_LSP_TIMEOUT", "1") + monkeypatch.setenv("AUTOFORM_MAX_LSP_REQUEST_SECONDS", "1") + + with pytest.raises(ValueError, match="REPL request and cleanup limits"): LeanRuntimeConfig.from_environment() From 855dd70d38f7bfdf04cbf05794e5de12f744a303 Mon Sep 17 00:00:00 2001 From: Jack McCarthy Date: Sat, 26 Sep 2026 14:23:53 -0400 Subject: [PATCH 2/5] Validate disposable REPL imports as Lean parses the header A disposable call sends its whole header to Lean with no environment, but the allowlist read it line by line. Lean skips comments before and between imports and accepts several imports on one line, so a header such as a block comment followed by an import, or two imports on one line, reached Lean without its later modules being checked. The disposable path now scans the header the way Lean does, reading every import across comments and lines, and rejects header syntax the scanner does not model (module, prelude, public/meta/private import, import all) so validation fails closed. --- servers/repl/core.py | 116 ++++++++++++++++++++++++++++++- tests/test_repl_core_protocol.py | 64 +++++++++++++++++ 2 files changed, 179 insertions(+), 1 deletion(-) diff --git a/servers/repl/core.py b/servers/repl/core.py index 2f8cd80b..5fb7dd15 100644 --- a/servers/repl/core.py +++ b/servers/repl/core.py @@ -290,6 +290,113 @@ def _validate_command_response( return environment, messages +# Keywords valid only at the start of a file, and words that modify an import. +_HEADER_ONLY_KEYWORDS = ("prelude", "module") +_IMPORT_MODIFIERS = ("public", "private", "meta") + + +def _header_import_modules(code: str) -> list[str]: + """Return every module the Lean header imports, as Lean would parse it. + + Unlike ``_split_imports_and_body``, this skips line and nested block comments + and reads several imports on one line, because a request sent without an + environment has its whole header processed by Lean. Header syntax this + scanner does not model raises ``ValueError`` so validation fails closed. + """ + position = 0 + length = len(code) + modules: list[str] = [] + + def skip_trivia() -> None: + nonlocal position + while position < length: + if code[position].isspace(): + position += 1 + elif code.startswith("--", position): + newline = code.find("\n", position) + position = length if newline < 0 else newline + 1 + elif code.startswith("/-", position): + depth = 0 + while position < length: + if code.startswith("/-", position): + depth += 1 + position += 2 + elif code.startswith("-/", position): + depth -= 1 + position += 2 + if depth == 0: + break + else: + position += 1 + else: + return + + def keyword_at(word: str) -> bool: + if not code.startswith(word, position): + return False + end = position + len(word) + return ( + end >= length + or code[end].isspace() + or code[end] == "«" + or code.startswith("--", end) + or code.startswith("/-", end) + ) + + def import_modifier_at() -> str | None: + """Return a modifier word only when a chain of them leads to import.""" + nonlocal position + start = position + first: str | None = None + try: + while True: + word = next( + (w for w in _IMPORT_MODIFIERS if keyword_at(w)), + None, + ) + if word is None: + return first if first is not None and keyword_at("import") else None + first = first or word + position += len(word) + skip_trivia() + finally: + position = start + + while True: + skip_trivia() + for keyword in _HEADER_ONLY_KEYWORDS: + if keyword_at(keyword): + raise ValueError(f"Unsupported Lean header keyword {keyword!r}.") + modifier = import_modifier_at() + if modifier is not None: + raise ValueError(f"Unsupported Lean header keyword {modifier!r}.") + if not keyword_at("import"): + return modules + position += len("import") + skip_trivia() + start = position + while position < length: + if code[position] == "«": + closing = code.find("»", position + 1) + if closing < 0: + raise ValueError("Unterminated module name in Lean header.") + position = closing + 1 + elif ( + code[position].isspace() + or code.startswith("--", position) + or code.startswith("/-", position) + ): + break + else: + position += 1 + module = code[start:position] + if not module: + raise ValueError("Missing module name after import.") + if module == "all": + raise ValueError("Unsupported Lean header keyword 'all'.") + modules.append(module) + + def _split_imports_and_body(code: str) -> tuple[list[str], str, int]: """Split Lean code into import statements and body. @@ -780,8 +887,15 @@ def remaining() -> float: self.config.validate_imports and self._allowed_import_roots is not None ): + # The whole header reaches Lean here, so validate it as Lean + # parses it rather than line by line. + try: + header_modules = _header_import_modules(code) + except ValueError as error: + result = {"repl_error": f"Rejected Lean header: {error}"} + header_modules = [] submitted_roots = { - statement.split(".")[0] for statement in imports + module.split(".")[0] for module in header_modules } disallowed = submitted_roots - self._allowed_import_roots if disallowed: diff --git a/tests/test_repl_core_protocol.py b/tests/test_repl_core_protocol.py index 09ca1637..e5257281 100644 --- a/tests/test_repl_core_protocol.py +++ b/tests/test_repl_core_protocol.py @@ -906,6 +906,70 @@ def close(*, deadline=None): assert repl.process is None + +@pytest.mark.parametrize( + ("code", "expected"), + [ + ("import Mathlib\n#check Nat", ["Mathlib"]), + ("/- note -/\nimport Unsafe\n", ["Unsafe"]), + ("import Mathlib.Tactic import Unsafe\n", ["Mathlib.Tactic", "Unsafe"]), + ("/- a /- nested -/ b -/ import Unsafe", ["Unsafe"]), + ("-- c\nimport Mathlib -- t\nimport Aesop/- x -/import Unsafe", ["Mathlib", "Aesop", "Unsafe"]), + ("private theorem t : True := trivial\nimport Unsafe", []), + ("import Mathlib\nprivate def f := 1", ["Mathlib"]), + ("theorem importUnsafe : True := trivial", []), + ], +) +def test_header_scan_reads_imports_as_lean_parses_them(code, expected): + assert repl_core._header_import_modules(code) == expected + + +@pytest.mark.parametrize( + "code", + [ + "public import Unsafe", + "meta import Unsafe", + "/- c -/ public /- d -/ import Unsafe", + "module\nimport Unsafe", + "prelude\nimport Unsafe", + "import all Unsafe", + "import «Unterminated", + ], +) +def test_header_scan_rejects_header_syntax_it_does_not_model(code): + with pytest.raises(ValueError): + repl_core._header_import_modules(code) + + +@pytest.mark.parametrize( + ("code", "expected_error"), + [ + ("/- note -/\nimport Unsafe\n#check Nat", "Disallowed imports: Unsafe"), + ("import Mathlib.Tactic import Unsafe\n#check Nat", "Disallowed imports: Unsafe"), + ("public import Unsafe\n#check Nat", "Rejected Lean header"), + ], +) +def test_run_disposable_rejects_imports_hidden_from_a_line_scan( + monkeypatch, code, expected_error +): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + allowed_imports=frozenset({"Mathlib"}), + warmup_imports=frozenset(), + ) + ) + monkeypatch.setattr( + repl, + "start", + lambda *args, **kwargs: pytest.fail("invalid input must not start Lean"), + ) + + response = repl.run_disposable(code) + + assert expected_error in response["repl_error"] + assert repl.process is None + + class _PipeProcess: def __init__(self, stack: ExitStack, stdout_chunks: list[bytes], stderr: bytes = b""): stdin_read, stdin_write = os.pipe() From 3a5351b684fe4cfbb7f09324e4e75324ee084e19 Mon Sep 17 00:00:00 2001 From: Jack McCarthy Date: Sat, 26 Sep 2026 15:04:49 -0400 Subject: [PATCH 3/5] Check disposable REPL imports with Lean's own header parser Replace the hand-written Python header scanner with `lean --stdin --deps`, run under `lake env` so it resolves imports against the worker's search path. Each resolved .olean path is mapped back to its module name and its root checked against the allowlist. A header Lean rejects, an unresolvable import, a path outside the search path, or a timeout rejects the run. The check now runs on the exact text sent to Lean, after warmup imports are prepended, and no longer needs special cases for `module`, `prelude`, or import modifiers. Adversarial header cases move to the real-REPL job so they run against the pinned Lean. --- servers/repl/core.py | 187 ++++++++++++------------------- tests/test_real_repl.py | 35 ++++++ tests/test_repl_core_protocol.py | 114 +++++++++++++------ 3 files changed, 190 insertions(+), 146 deletions(-) diff --git a/servers/repl/core.py b/servers/repl/core.py index 5fb7dd15..70fd4a64 100644 --- a/servers/repl/core.py +++ b/servers/repl/core.py @@ -290,111 +290,59 @@ def _validate_command_response( return environment, messages -# Keywords valid only at the start of a file, and words that modify an import. -_HEADER_ONLY_KEYWORDS = ("prelude", "module") -_IMPORT_MODIFIERS = ("public", "private", "meta") - - -def _header_import_modules(code: str) -> list[str]: - """Return every module the Lean header imports, as Lean would parse it. - - Unlike ``_split_imports_and_body``, this skips line and nested block comments - and reads several imports on one line, because a request sent without an - environment has its whole header processed by Lean. Header syntax this - scanner does not model raises ``ValueError`` so validation fails closed. +def _lean_header_modules( + command: list[str], + code: str, + *, + cwd: str | None, + env: dict[str, str], + deadline: float, +) -> list[str]: + """Return every module the Lean header of ``code`` imports, read by Lean itself. + + ``command`` prints the Lean search path on its first line, then runs + ``lean --stdin --deps``, which parses the header with Lean's own parser and + prints one resolved ``.olean`` path per import. Each path is mapped back to + its module name. A header Lean rejects, an import it cannot resolve, or a + path outside the search path raises ``ValueError`` so validation fails closed. """ - position = 0 - length = len(code) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("timed out before checking the Lean header") + process = subprocess.Popen( + command, + cwd=cwd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + start_new_session=True, + ) + try: + stdout, stderr = process.communicate(code.encode(), timeout=remaining) + except subprocess.TimeoutExpired: + _kill_subprocesses(process, process.pid) + raise TimeoutError("timed out checking the Lean header") from None + except BaseException: + _kill_subprocesses(process, process.pid) + raise + if process.returncode != 0: + detail = stderr.decode(errors="replace").strip().splitlines() + raise ValueError(detail[0] if detail else f"exit status {process.returncode}") + + search_path, *paths = stdout.decode(errors="replace").splitlines() or [""] + roots = [os.path.normpath(entry) for entry in search_path.split(os.pathsep) if entry] modules: list[str] = [] - - def skip_trivia() -> None: - nonlocal position - while position < length: - if code[position].isspace(): - position += 1 - elif code.startswith("--", position): - newline = code.find("\n", position) - position = length if newline < 0 else newline + 1 - elif code.startswith("/-", position): - depth = 0 - while position < length: - if code.startswith("/-", position): - depth += 1 - position += 2 - elif code.startswith("-/", position): - depth -= 1 - position += 2 - if depth == 0: - break - else: - position += 1 - else: - return - - def keyword_at(word: str) -> bool: - if not code.startswith(word, position): - return False - end = position + len(word) - return ( - end >= length - or code[end].isspace() - or code[end] == "«" - or code.startswith("--", end) - or code.startswith("/-", end) - ) - - def import_modifier_at() -> str | None: - """Return a modifier word only when a chain of them leads to import.""" - nonlocal position - start = position - first: str | None = None - try: - while True: - word = next( - (w for w in _IMPORT_MODIFIERS if keyword_at(w)), - None, - ) - if word is None: - return first if first is not None and keyword_at("import") else None - first = first or word - position += len(word) - skip_trivia() - finally: - position = start - - while True: - skip_trivia() - for keyword in _HEADER_ONLY_KEYWORDS: - if keyword_at(keyword): - raise ValueError(f"Unsupported Lean header keyword {keyword!r}.") - modifier = import_modifier_at() - if modifier is not None: - raise ValueError(f"Unsupported Lean header keyword {modifier!r}.") - if not keyword_at("import"): - return modules - position += len("import") - skip_trivia() - start = position - while position < length: - if code[position] == "«": - closing = code.find("»", position + 1) - if closing < 0: - raise ValueError("Unterminated module name in Lean header.") - position = closing + 1 - elif ( - code[position].isspace() - or code.startswith("--", position) - or code.startswith("/-", position) - ): - break - else: - position += 1 - module = code[start:position] - if not module: - raise ValueError("Missing module name after import.") - if module == "all": - raise ValueError("Unsupported Lean header keyword 'all'.") - modules.append(module) + for path in paths: + path = os.path.normpath(path) + root = next((r for r in roots if path.startswith(r + os.sep)), None) + if root is None or not path.endswith(".olean"): + raise ValueError(f"import resolved outside the Lean search path: {path}") + module = path[len(root) + 1 : -len(".olean")].replace(os.sep, ".") + # Every file imports Init implicitly, so importing it grants nothing. + if module != "Init": + modules.append(module) + return modules def _split_imports_and_body(code: str) -> tuple[list[str], str, int]: @@ -445,6 +393,12 @@ class LeanReplConfig: warmup_imports: frozenset[str] = WARMUP_IMPORTS repl_command: list[str] = field(default_factory=lambda: ["lake", "exe", "repl"]) + # Prints LEAN_PATH, then Lean's own resolution of the submitted header. + header_deps_command: list[str] = field( + default_factory=lambda: [ + "lake", "env", "sh", "-c", 'printenv LEAN_PATH && exec lean --stdin --deps' + ] + ) # stdout is capped per response. stderr has no protocol framing, so its # ceiling applies to the entire process generation and resets on restart. @@ -883,14 +837,28 @@ def remaining() -> float: try: self.close(deadline=deadline) imports, _, _ = _split_imports_and_body(code) + added_imports = tuple( + root + for root in sorted(self.config.warmup_imports) + if root not in imports + ) + prefix = "\n".join(f"import {root}" for root in added_imports) + command = f"{prefix}\n{code}" if prefix else code if ( self.config.validate_imports and self._allowed_import_roots is not None ): - # The whole header reaches Lean here, so validate it as Lean - # parses it rather than line by line. + # The whole header reaches Lean here, so let Lean parse it. + env = _inherit_clean_env() + env.update(self.config.env) try: - header_modules = _header_import_modules(code) + header_modules = _lean_header_modules( + self.config.header_deps_command, + command, + cwd=self.cwd, + env=env, + deadline=deadline, + ) except ValueError as error: result = {"repl_error": f"Rejected Lean header: {error}"} header_modules = [] @@ -907,13 +875,6 @@ def remaining() -> float: ) } if result is None: - added_imports = tuple( - root - for root in sorted(self.config.warmup_imports) - if root not in imports - ) - prefix = "\n".join(f"import {root}" for root in added_imports) - command = f"{prefix}\n{code}" if prefix else code self.start(deadline=deadline, warmup_imports=()) response = self._run( code=command, diff --git a/tests/test_real_repl.py b/tests/test_real_repl.py index 81dc946d..bfad33cc 100644 --- a/tests/test_real_repl.py +++ b/tests/test_real_repl.py @@ -36,3 +36,38 @@ def test_disposable_call_matches_the_pinned_repl_protocol(): assert "env" not in response assert all("proofState" not in sorry for sorry in response["sorries"]) assert repl.is_clean() + + +@pytest.mark.skipif( + os.environ.get("AUTOFORM_RUN_REAL_REPL_TESTS") != "1", + reason="set AUTOFORM_RUN_REAL_REPL_TESTS=1 to run the pinned REPL integration", +) +@pytest.mark.parametrize( + ("warmup", "code", "expected_error"), + [ + ((), "/- note -/\nimport Init.Data\n#check Nat", "Disallowed imports: Init"), + ((), "import REPL import Init.Data\n#check Nat", "Disallowed imports: Init"), + ((), "module\npublic import Init.Data\n", "Disallowed imports: Init"), + (("REPL",), "/- note -/ import Init.Data\n#check Nat", "Disallowed imports: Init"), + ((), "import NotAllowlisted.Mod\n", "Rejected Lean header"), + ((), "import «REPL\n", "Rejected Lean header"), + ((), "import REPL\ntheorem autoform_header_probe : True := trivial", None), + ], +) +def test_disposable_imports_are_checked_by_lean_itself(warmup, code, expected_error): + repl = LeanRepl( + LeanReplConfig( + cwd=str(REPL_FIXTURE), + repl_command=["lake", "exe", "repl"], + allowed_imports=frozenset({"REPL"}), + warmup_imports=frozenset(warmup), + ) + ) + + response = repl.run_disposable(code, timeout=180) + + if expected_error is None: + assert "repl_error" not in response + else: + assert expected_error in response["repl_error"] + assert repl.is_clean() diff --git a/tests/test_repl_core_protocol.py b/tests/test_repl_core_protocol.py index e5257281..a1c3aad0 100644 --- a/tests/test_repl_core_protocol.py +++ b/tests/test_repl_core_protocol.py @@ -883,6 +883,7 @@ def test_run_disposable_closes_before_rejecting_an_import(monkeypatch): repl_core.LeanReplConfig( allowed_imports=frozenset({"Mathlib"}), warmup_imports=frozenset(), + header_deps_command=_fake_header_deps("/lib\n/lib/Unsafe.olean\n"), ) ) repl.process = object() @@ -907,55 +908,82 @@ def close(*, deadline=None): -@pytest.mark.parametrize( - ("code", "expected"), - [ - ("import Mathlib\n#check Nat", ["Mathlib"]), - ("/- note -/\nimport Unsafe\n", ["Unsafe"]), - ("import Mathlib.Tactic import Unsafe\n", ["Mathlib.Tactic", "Unsafe"]), - ("/- a /- nested -/ b -/ import Unsafe", ["Unsafe"]), - ("-- c\nimport Mathlib -- t\nimport Aesop/- x -/import Unsafe", ["Mathlib", "Aesop", "Unsafe"]), - ("private theorem t : True := trivial\nimport Unsafe", []), - ("import Mathlib\nprivate def f := 1", ["Mathlib"]), - ("theorem importUnsafe : True := trivial", []), - ], -) -def test_header_scan_reads_imports_as_lean_parses_them(code, expected): - assert repl_core._header_import_modules(code) == expected +def _fake_header_deps(stdout: str, returncode: int = 0, stderr: str = "") -> list[str]: + """Stand in for ``lake env ... lean --stdin --deps`` with fixed output.""" + script = ( + "import sys; sys.stdin.read(); " + f"sys.stdout.write({stdout!r}); sys.stderr.write({stderr!r}); " + f"sys.exit({returncode})" + ) + return [sys.executable, "-c", script] + + +def _header_modules(command: list[str], deadline: float | None = None) -> list[str]: + return repl_core._lean_header_modules( + command, + "import Mathlib", + cwd=None, + env=dict(os.environ), + deadline=time.monotonic() + 10 if deadline is None else deadline, + ) + + +def test_header_check_maps_resolved_paths_back_to_module_names(): + command = _fake_header_deps( + "/pkg/lib:/toolchain/lib\n" + "/toolchain/lib/Init.olean\n" + "/toolchain/lib/Init.olean\n" + "/pkg/lib/Mathlib/Tactic.olean\n" + "/toolchain/lib/Init/Data.olean\n" + ) + + assert _header_modules(command) == ["Mathlib.Tactic", "Init.Data"] @pytest.mark.parametrize( - "code", + ("command", "message"), [ - "public import Unsafe", - "meta import Unsafe", - "/- c -/ public /- d -/ import Unsafe", - "module\nimport Unsafe", - "prelude\nimport Unsafe", - "import all Unsafe", - "import «Unterminated", + ( + _fake_header_deps("", 1, "unknown module prefix 'Unsafe'\nmore"), + "unknown module prefix 'Unsafe'", + ), + (_fake_header_deps("", 3), "exit status 3"), + (_fake_header_deps("/pkg/lib\n/other/Unsafe.olean\n"), "outside the Lean search path"), + (_fake_header_deps("/pkg/lib\n/pkg/lib/Unsafe.ilean\n"), "outside the Lean search path"), ], ) -def test_header_scan_rejects_header_syntax_it_does_not_model(code): - with pytest.raises(ValueError): - repl_core._header_import_modules(code) +def test_header_check_fails_closed(command, message): + with pytest.raises(ValueError, match=message): + _header_modules(command) + + +def test_header_check_kills_a_command_that_outlives_the_deadline(): + command = [sys.executable, "-c", "import time; time.sleep(30)"] + started = time.monotonic() + + with pytest.raises(TimeoutError): + _header_modules(command, deadline=started + 0.2) + + assert time.monotonic() - started < 5 @pytest.mark.parametrize( - ("code", "expected_error"), + ("deps_output", "returncode", "expected_error"), [ - ("/- note -/\nimport Unsafe\n#check Nat", "Disallowed imports: Unsafe"), - ("import Mathlib.Tactic import Unsafe\n#check Nat", "Disallowed imports: Unsafe"), - ("public import Unsafe\n#check Nat", "Rejected Lean header"), + ("/lib\n/lib/Mathlib.olean\n/lib/Unsafe.olean\n", 0, "Disallowed imports: Unsafe"), + ("", 1, "Rejected Lean header: unknown module"), ], ) -def test_run_disposable_rejects_imports_hidden_from_a_line_scan( - monkeypatch, code, expected_error +def test_run_disposable_rejects_what_lean_reports_before_starting( + monkeypatch, deps_output, returncode, expected_error ): repl = repl_core.LeanRepl( repl_core.LeanReplConfig( allowed_imports=frozenset({"Mathlib"}), warmup_imports=frozenset(), + header_deps_command=_fake_header_deps( + deps_output, returncode, "unknown module prefix 'Unsafe'" + ), ) ) monkeypatch.setattr( @@ -964,12 +992,32 @@ def test_run_disposable_rejects_imports_hidden_from_a_line_scan( lambda *args, **kwargs: pytest.fail("invalid input must not start Lean"), ) - response = repl.run_disposable(code) + response = repl.run_disposable("import Mathlib\n#check Nat") assert expected_error in response["repl_error"] assert repl.process is None +def test_run_disposable_checks_the_header_after_adding_warmup_imports(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + allowed_imports=frozenset({"Mathlib"}), + warmup_imports=frozenset({"Mathlib"}), + ) + ) + checked = [] + + def header_modules(command, code, **kwargs): + checked.append(code) + raise ValueError("stop here") + + monkeypatch.setattr(repl_core, "_lean_header_modules", header_modules) + + repl.run_disposable("/- note -/ import Unsafe\n#check Nat") + + assert checked == ["import Mathlib\n/- note -/ import Unsafe\n#check Nat"] + + class _PipeProcess: def __init__(self, stack: ExitStack, stdout_chunks: list[bytes], stderr: bytes = b""): stdin_read, stdin_write = os.pipe() From e3f40ab36c67464b1cf0d41927ba7d0026dc0cb0 Mon Sep 17 00:00:00 2001 From: Jack McCarthy Date: Sat, 26 Sep 2026 15:25:13 -0400 Subject: [PATCH 4/5] Read header imports from lean --deps-json and keep pool cleanup deadlines exact MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `lean --deps` reports resolved .olean paths, and mapping them back to module names misreads a quoted atomic name such as «Mathlib.X» as root Mathlib. `lean --deps-json` reports each module name as Lean spells it, needs no search-path mapping, and reports header parse errors in its output, which now reject the run. Its output shape is unchanged from Lean 4.26 to 4.34. LeanRepl.close_with_deadline no longer widens the pool's shared deadline per worker, so shutting down several workers cannot outlast the cleanup budget. close() still reserves cleanup time after an expired request deadline. Also remove the unused startup_stagger setting, document that import filtering is not a security sandbox, and add a real-REPL test that the same declaration succeeds on two consecutive pool calls. --- servers/README.md | 14 ++++--- servers/repl/core.py | 52 +++++++++++++----------- servers/repl/pool.py | 2 - tests/test_real_repl.py | 29 ++++++++++++- tests/test_repl_core_protocol.py | 67 ++++++++++++++++++++++--------- tests/test_repl_pool_lifecycle.py | 28 ++++++------- 6 files changed, 126 insertions(+), 66 deletions(-) diff --git a/servers/README.md b/servers/README.md index eabf24a5..46c7ab46 100644 --- a/servers/README.md +++ b/servers/README.md @@ -15,11 +15,15 @@ project. Every public REPL call starts a fresh child and reaps it before returning, so environments, proof states, and stream contents cannot cross independent successful calls. If cleanup cannot be confirmed, Autoform returns an explicit no-replay error, quarantines that project pool, and blocks its -replacement until cleanup succeeds. LSP sessions remain resident because their -protocol is explicitly stateful. Closing the session that started the runtime -does not stop it; after a crash, the next tool call starts it again. Runtime -sockets include a code fingerprint, so an in-place upgrade gracefully replaces -the older build. +replacement until cleanup succeeds. Before a child starts, Lean itself parses +the submitted header (`lean --deps-json`) and imports outside the allowed roots +(`Mathlib`, `Aesop`, `Batteries`, `LeanSearchClient` by default) are rejected. +This keeps calls on known libraries; it is not a security sandbox, because the +submitted Lean code can still run arbitrary `IO`. LSP sessions remain resident +because their protocol is explicitly stateful. Closing the session that started +the runtime does not stop it; after a crash, the next tool call starts it again. +Runtime sockets include a code fingerprint, so an in-place upgrade gracefully +replaces the older build. Lean subprocesses remain lazy. A REPL call stays pending while its fresh child starts, and the first LSP call stays pending while its session starts, so no diff --git a/servers/repl/core.py b/servers/repl/core.py index 70fd4a64..a0fff95f 100644 --- a/servers/repl/core.py +++ b/servers/repl/core.py @@ -300,11 +300,11 @@ def _lean_header_modules( ) -> list[str]: """Return every module the Lean header of ``code`` imports, read by Lean itself. - ``command`` prints the Lean search path on its first line, then runs - ``lean --stdin --deps``, which parses the header with Lean's own parser and - prints one resolved ``.olean`` path per import. Each path is mapped back to - its module name. A header Lean rejects, an import it cannot resolve, or a - path outside the search path raises ``ValueError`` so validation fails closed. + ``command`` runs ``lean --deps-json /dev/stdin``, which parses the header with + Lean's own parser and reports each import's module name as Lean spells it, + so a quoted name such as ``«Mathlib.X»`` keeps its quotes and root. A header + Lean rejects, a failed command, or output in an unknown shape raises + ``ValueError`` so validation fails closed. """ remaining = deadline - time.monotonic() if remaining <= 0: @@ -330,19 +330,18 @@ def _lean_header_modules( detail = stderr.decode(errors="replace").strip().splitlines() raise ValueError(detail[0] if detail else f"exit status {process.returncode}") - search_path, *paths = stdout.decode(errors="replace").splitlines() or [""] - roots = [os.path.normpath(entry) for entry in search_path.split(os.pathsep) if entry] - modules: list[str] = [] - for path in paths: - path = os.path.normpath(path) - root = next((r for r in roots if path.startswith(r + os.sep)), None) - if root is None or not path.endswith(".olean"): - raise ValueError(f"import resolved outside the Lean search path: {path}") - module = path[len(root) + 1 : -len(".olean")].replace(os.sep, ".") - # Every file imports Init implicitly, so importing it grants nothing. - if module != "Init": - modules.append(module) - return modules + try: + (entry,) = json.loads(stdout)["imports"] + errors = entry["errors"] + modules = [] if errors else [item["module"] for item in entry["result"]["imports"]] + except (ValueError, TypeError, KeyError): + raise ValueError("unrecognized output from lean --deps-json") from None + if errors: + raise ValueError(str(errors[0])) + if not all(isinstance(module, str) and module for module in modules): + raise ValueError("unrecognized output from lean --deps-json") + # Every file imports Init implicitly, so importing it grants nothing. + return [module for module in modules if module != "Init"] def _split_imports_and_body(code: str) -> tuple[list[str], str, int]: @@ -393,11 +392,9 @@ class LeanReplConfig: warmup_imports: frozenset[str] = WARMUP_IMPORTS repl_command: list[str] = field(default_factory=lambda: ["lake", "exe", "repl"]) - # Prints LEAN_PATH, then Lean's own resolution of the submitted header. + # Reports the imports of the submitted header, parsed by Lean itself. header_deps_command: list[str] = field( - default_factory=lambda: [ - "lake", "env", "sh", "-c", 'printenv LEAN_PATH && exec lean --stdin --deps' - ] + default_factory=lambda: ["lake", "env", "lean", "--deps-json", "/dev/stdin"] ) # stdout is capped per response. stderr has no protocol framing, so its @@ -748,6 +745,9 @@ def close(self, *, deadline: float | None = None) -> None: """ if deadline is not None: deadline = max(deadline, time.monotonic() + DEFAULT_REPL_CLEANUP_SECONDS) + self._close(deadline) + + def _close(self, deadline: float | None) -> None: process = self.process process_group_id = self._process_group_id try: @@ -774,8 +774,12 @@ def close(self, *, deadline: float | None = None) -> None: self._stderr_tail.clear() def close_with_deadline(self, deadline: float) -> None: - """Close using an absolute deadline shared by a pool shutdown.""" - self.close(deadline=deadline) + """Close within an absolute deadline shared by a pool shutdown. + + Unlike close(), this never extends the deadline, so closing several + workers in turn cannot outlast the pool's cleanup budget. + """ + self._close(deadline) def is_clean(self) -> bool: """Return whether this wrapper owns no live or unreaped process group.""" diff --git a/servers/repl/pool.py b/servers/repl/pool.py index 8de493c5..24ae5a16 100644 --- a/servers/repl/pool.py +++ b/servers/repl/pool.py @@ -20,7 +20,6 @@ DEFAULT_PORT = 8990 DEFAULT_RAM_FRACTION = 0.5 -DEFAULT_STARTUP_STAGGER_SECONDS = 2.0 DEFAULT_POOL_CLEANUP_SECONDS = DEFAULT_REPL_CLEANUP_SECONDS @@ -37,7 +36,6 @@ class LeanReplPoolConfig(LeanReplConfig): """Configuration for a pool of Lean REPL instances.""" num_repls: int | None = None - startup_stagger: float = DEFAULT_STARTUP_STAGGER_SECONDS def __post_init__(self) -> None: if self.num_repls is None: diff --git a/tests/test_real_repl.py b/tests/test_real_repl.py index bfad33cc..f6bb2d78 100644 --- a/tests/test_real_repl.py +++ b/tests/test_real_repl.py @@ -8,6 +8,7 @@ import pytest from servers.repl.core import LeanRepl, LeanReplConfig +from servers.repl.pool import LeanReplPool, LeanReplPoolConfig REPL_FIXTURE = Path(__file__).parent / "fixtures" / "repl-smoke" @@ -49,7 +50,8 @@ def test_disposable_call_matches_the_pinned_repl_protocol(): ((), "import REPL import Init.Data\n#check Nat", "Disallowed imports: Init"), ((), "module\npublic import Init.Data\n", "Disallowed imports: Init"), (("REPL",), "/- note -/ import Init.Data\n#check Nat", "Disallowed imports: Init"), - ((), "import NotAllowlisted.Mod\n", "Rejected Lean header"), + ((), "import NotAllowlisted.Mod\n", "Disallowed imports: NotAllowlisted"), + ((), "import «REPL.X»\n", "Disallowed imports: «REPL"), ((), "import «REPL\n", "Rejected Lean header"), ((), "import REPL\ntheorem autoform_header_probe : True := trivial", None), ], @@ -71,3 +73,28 @@ def test_disposable_imports_are_checked_by_lean_itself(warmup, code, expected_er else: assert expected_error in response["repl_error"] assert repl.is_clean() + + +@pytest.mark.skipif( + os.environ.get("AUTOFORM_RUN_REAL_REPL_TESTS") != "1", + reason="set AUTOFORM_RUN_REAL_REPL_TESTS=1 to run the pinned REPL integration", +) +def test_pool_calls_do_not_share_lean_state(): + pool = LeanReplPool( + LeanReplPoolConfig( + cwd=str(REPL_FIXTURE), + repl_command=["lake", "exe", "repl"], + allowed_imports=frozenset({"REPL"}), + warmup_imports=frozenset(), + num_repls=1, + ) + ) + declaration = "theorem autoform_isolation_probe : True := trivial" + try: + responses = [pool.run(declaration, timeout=180) for _ in range(2)] + finally: + pool.shutdown() + + for response in responses: + assert "repl_error" not in response + assert not any(m["severity"] == "error" for m in response.get("messages", [])) diff --git a/tests/test_repl_core_protocol.py b/tests/test_repl_core_protocol.py index a1c3aad0..a0427474 100644 --- a/tests/test_repl_core_protocol.py +++ b/tests/test_repl_core_protocol.py @@ -883,7 +883,7 @@ def test_run_disposable_closes_before_rejecting_an_import(monkeypatch): repl_core.LeanReplConfig( allowed_imports=frozenset({"Mathlib"}), warmup_imports=frozenset(), - header_deps_command=_fake_header_deps("/lib\n/lib/Unsafe.olean\n"), + header_deps_command=_fake_header_deps(_deps_json("Unsafe")), ) ) repl.process = object() @@ -909,7 +909,7 @@ def close(*, deadline=None): def _fake_header_deps(stdout: str, returncode: int = 0, stderr: str = "") -> list[str]: - """Stand in for ``lake env ... lean --stdin --deps`` with fixed output.""" + """Stand in for ``lake env lean --deps-json /dev/stdin`` with fixed output.""" script = ( "import sys; sys.stdin.read(); " f"sys.stdout.write({stdout!r}); sys.stderr.write({stderr!r}); " @@ -918,6 +918,16 @@ def _fake_header_deps(stdout: str, returncode: int = 0, stderr: str = "") -> lis return [sys.executable, "-c", script] +def _deps_json(*modules: str, errors: tuple[str, ...] = ()) -> str: + entry: dict = {"errors": list(errors)} + if not errors: + entry["result"] = { + "imports": [{"module": module, "importAll": False} for module in modules], + "isModule": False, + } + return json.dumps({"imports": [entry]}) + + def _header_modules(command: list[str], deadline: float | None = None) -> list[str]: return repl_core._lean_header_modules( command, @@ -928,28 +938,27 @@ def _header_modules(command: list[str], deadline: float | None = None) -> list[s ) -def test_header_check_maps_resolved_paths_back_to_module_names(): +def test_header_check_returns_the_modules_lean_reports(): command = _fake_header_deps( - "/pkg/lib:/toolchain/lib\n" - "/toolchain/lib/Init.olean\n" - "/toolchain/lib/Init.olean\n" - "/pkg/lib/Mathlib/Tactic.olean\n" - "/toolchain/lib/Init/Data.olean\n" + _deps_json("Init", "Init", "Mathlib.Tactic", "«Mathlib.X»", "Init.Data") ) - assert _header_modules(command) == ["Mathlib.Tactic", "Init.Data"] + assert _header_modules(command) == ["Mathlib.Tactic", "«Mathlib.X»", "Init.Data"] @pytest.mark.parametrize( ("command", "message"), [ ( - _fake_header_deps("", 1, "unknown module prefix 'Unsafe'\nmore"), - "unknown module prefix 'Unsafe'", + _fake_header_deps(_deps_json(errors=(":2:0: unterminated identifier escape",))), + "unterminated identifier escape", ), + (_fake_header_deps("", 1, "lake: unknown\nmore"), "lake: unknown"), (_fake_header_deps("", 3), "exit status 3"), - (_fake_header_deps("/pkg/lib\n/other/Unsafe.olean\n"), "outside the Lean search path"), - (_fake_header_deps("/pkg/lib\n/pkg/lib/Unsafe.ilean\n"), "outside the Lean search path"), + (_fake_header_deps("not json"), "unrecognized output"), + (_fake_header_deps('{"imports": []}'), "unrecognized output"), + (_fake_header_deps('{"imports": [{"errors": [], "result": {}}]}'), "unrecognized output"), + (_fake_header_deps(_deps_json("")), "unrecognized output"), ], ) def test_header_check_fails_closed(command, message): @@ -968,22 +977,21 @@ def test_header_check_kills_a_command_that_outlives_the_deadline(): @pytest.mark.parametrize( - ("deps_output", "returncode", "expected_error"), + ("deps_output", "expected_error"), [ - ("/lib\n/lib/Mathlib.olean\n/lib/Unsafe.olean\n", 0, "Disallowed imports: Unsafe"), - ("", 1, "Rejected Lean header: unknown module"), + (_deps_json("Mathlib", "Unsafe.Mod"), "Disallowed imports: Unsafe"), + (_deps_json("«Mathlib.X»"), "Disallowed imports: «Mathlib"), + (_deps_json(errors=("bad header",)), "Rejected Lean header: bad header"), ], ) def test_run_disposable_rejects_what_lean_reports_before_starting( - monkeypatch, deps_output, returncode, expected_error + monkeypatch, deps_output, expected_error ): repl = repl_core.LeanRepl( repl_core.LeanReplConfig( allowed_imports=frozenset({"Mathlib"}), warmup_imports=frozenset(), - header_deps_command=_fake_header_deps( - deps_output, returncode, "unknown module prefix 'Unsafe'" - ), + header_deps_command=_fake_header_deps(deps_output), ) ) monkeypatch.setattr( @@ -2019,6 +2027,25 @@ def test_close_after_an_expired_deadline_still_reserves_cleanup_time(monkeypatch assert deadlines[0] >= before + repl_core.DEFAULT_REPL_CLEANUP_SECONDS +def test_close_with_deadline_keeps_a_shared_pool_deadline_exact(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig(warmup_imports=frozenset(), validate_imports=False) + ) + repl.process = object() + repl._process_group_id = 1234 + deadlines = [] + monkeypatch.setattr( + repl_core, + "_kill_subprocesses", + lambda process, process_group_id, deadline=None: deadlines.append(deadline), + ) + + shared_deadline = time.monotonic() - 5 + repl.close_with_deadline(shared_deadline) + + assert deadlines == [shared_deadline] + + def test_request_timeout_reaps_a_slow_to_exit_process_without_cleanup_failure( monkeypatch, ): diff --git a/tests/test_repl_pool_lifecycle.py b/tests/test_repl_pool_lifecycle.py index 1baf17dd..e3e08765 100644 --- a/tests/test_repl_pool_lifecycle.py +++ b/tests/test_repl_pool_lifecycle.py @@ -27,7 +27,7 @@ def close(self): self.closed = True monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) - config = repl_pool.LeanReplPoolConfig(num_repls=3, startup_stagger=0) + config = repl_pool.LeanReplPoolConfig(num_repls=3) pool = repl_pool.LeanReplPool(config) @@ -54,7 +54,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=2, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=2) ) pool.shutdown() @@ -79,7 +79,7 @@ def close_with_deadline(self, deadline): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=3, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=3) ) pool.shutdown() @@ -101,7 +101,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) borrowed = pool._idle.get_nowait() try: @@ -128,7 +128,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) try: with pytest.raises(OSError, match="stdout failed"): @@ -156,7 +156,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) try: @@ -186,7 +186,7 @@ def close_with_deadline(self, deadline): monkeypatch.setattr(repl_pool.time, "monotonic", lambda: now[0]) monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) try: @@ -217,7 +217,7 @@ def close(self): monkeypatch.setattr(repl_pool.time, "monotonic", lambda: now[0]) monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) try: @@ -249,7 +249,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) response = pool.run("#check Nat") @@ -286,7 +286,7 @@ def is_clean(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) try: @@ -311,7 +311,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) with pytest.raises(KeyboardInterrupt) as raised: @@ -340,7 +340,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) with pytest.raises(RuntimeError, match="cleanup failed"): @@ -378,7 +378,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) first = threading.Thread(target=pool.run, args=("first",)) first.start() @@ -435,7 +435,7 @@ def close(self): monkeypatch.setattr(repl_pool, "LeanRepl", FakeRepl) pool = repl_pool.LeanReplPool( - repl_pool.LeanReplPoolConfig(num_repls=1, startup_stagger=0) + repl_pool.LeanReplPoolConfig(num_repls=1) ) first = threading.Thread(target=pool.shutdown) From 55dd2cc1ffd9936d15623f76a3824b5f9af7c0b3 Mon Sep 17 00:00:00 2001 From: Jack McCarthy Date: Sat, 26 Sep 2026 17:34:31 -0400 Subject: [PATCH 5/5] Harden Lean header validation and pin REPL runtime --- .github/workflows/tests.yml | 3 +- README.md | 5 +- servers/README.md | 12 +- servers/lean_runtime.py | 2 +- servers/repl/core.py | 308 ++++++++++++++++-- skills/setup/SKILL.md | 6 +- .../assets/cabannes-thesis-project/README.md | 3 +- .../lake-manifest.json | 106 ++++++ .../cabannes-thesis-project/lakefile.toml | 5 + tests/fixtures/repl-smoke/Mathlib.lean | 1 + tests/fixtures/repl-smoke/lakefile.toml | 3 + tests/fixtures/repl-smoke/lean-toolchain | 2 +- tests/test_real_repl.py | 50 ++- tests/test_repl_core_protocol.py | 232 ++++++++++++- tests/test_shared_lean_runtime.py | 12 +- tests/test_skill_examples.py | 36 +- 16 files changed, 723 insertions(+), 63 deletions(-) create mode 100644 skills/setup/assets/cabannes-thesis-project/lake-manifest.json create mode 100644 tests/fixtures/repl-smoke/Mathlib.lean diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2387eb5d..8b4c214c 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -2,6 +2,7 @@ name: tests on: push: + branches: [main, execution] pull_request: permissions: @@ -38,7 +39,7 @@ jobs: with: auto-config: false build: true - build-args: repl + build-args: "Mathlib @repl/repl" lake-package-directory: tests/fixtures/repl-smoke - run: uv sync --extra dev --extra repl - run: uv run pytest -q tests/test_real_repl.py diff --git a/README.md b/README.md index 3f24013c..96d6b221 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,10 @@ manifest is included, but Muse installation is not covered here. ## Quick start Work from an existing Lean repository. First scaffold the blueprint and site -configuration from an Autoform checkout: +configuration from an Autoform checkout. The project must declare +`leanprover-community/repl` at an immutable revision compatible with its Lean +toolchain; verify it first with `lake build @repl/repl` (the Setup skill selects +and checks this pin): ```bash uv run autoform init /path/to/lean-project \ diff --git a/servers/README.md b/servers/README.md index 46c7ab46..66823cda 100644 --- a/servers/README.md +++ b/servers/README.md @@ -25,6 +25,12 @@ the runtime does not stop it; after a crash, the next tool call starts it again. Runtime sockets include a code fingerprint, so an in-place upgrade gracefully replaces the older build. +Each consumer Lake project must declare `leanprover-community/repl` at an +immutable revision compatible with its Lean toolchain. Autoform invokes the +qualified target with `lake exe @repl/repl`, so a same-named executable in the +root project cannot shadow the pinned dependency. `lake build @repl/repl` +checks this contract when the toolchain or dependency revision changes. + Lean subprocesses remain lazy. A REPL call stays pending while its fresh child starts, and the first LSP call stays pending while its session starts, so no `/repl-start`, `/lsp-start`, or model-side sleep is needed. Idle project slots @@ -47,9 +53,9 @@ request frame is dispatched. Once the final frame delimiter may have reached Lean, replay could execute the command twice, so Autoform retires the process and reports that the outcome is unknown instead of retrying. The REPL per-call timeout starts before the shared daemon is connected or -started, then covers project admission, fresh child startup, idle-worker wait, -and Lean execution. Verified process cleanup and response delivery get a -separate bounded grace period before the RPC returns. +started, then covers project admission, header validation, fresh child startup, +idle-worker wait, and Lean execution. Verified process cleanup and response +delivery get a separate bounded grace period before the RPC returns. `LEAN_REPL_CMD` is a trusted local command. Its descendants must remain in the dedicated process group Autoform creates; a command that deliberately detaches diff --git a/servers/lean_runtime.py b/servers/lean_runtime.py index 6d72dc38..1599ac8f 100644 --- a/servers/lean_runtime.py +++ b/servers/lean_runtime.py @@ -203,7 +203,7 @@ def from_environment(cls) -> "LeanRuntimeConfig": "AUTOFORM_REPL_TOTAL_WORKERS" ) repl_project_limit = min(max_projects, total_workers // workers_per_project) - repl_command = tuple(shlex.split(os.environ.get("LEAN_REPL_CMD", "lake exe repl"))) + repl_command = tuple(shlex.split(os.environ.get("LEAN_REPL_CMD", "lake exe @repl/repl"))) lsp_command = tuple(shlex.split(os.environ.get("LEAN_LSP_CMD", "lake serve"))) if not repl_command: raise ValueError("LEAN_REPL_CMD must not be empty") diff --git a/servers/repl/core.py b/servers/repl/core.py index a0fff95f..07017909 100644 --- a/servers/repl/core.py +++ b/servers/repl/core.py @@ -1,4 +1,4 @@ -"""Lean REPL backend: one session managing a ``lake exe repl`` subprocess. +"""Lean REPL backend: one session managing a ``lake exe @repl/repl`` subprocess. Provides LeanRepl with non-blocking I/O, a preloaded import environment, memory monitoring, automatic restart, and multi-snippet chaining. @@ -13,6 +13,7 @@ import select import signal import subprocess +import sys import threading import time from dataclasses import dataclass, field @@ -34,6 +35,11 @@ _STDERR_TAIL_BYTES = 200 _PUBLIC_DIAGNOSTIC_FIELDS = frozenset({"severity", "data", "pos", "endPos"}) _PUBLIC_SORRY_FIELDS = frozenset({"goal", "pos", "endPos"}) +_LEAN_HEADER_LAUNCHER = ( + "import os; " + "lean = os.path.join(os.environ['LEAN_SYSROOT'], 'bin', 'lean'); " + "os.execv(lean, [lean, '--deps-json', '/dev/stdin'])" +) # --------------------------------------------------------------------------- @@ -290,6 +296,197 @@ def _validate_command_response( return environment, messages +class _HeaderProcessCleanupError(RuntimeError): + """Header validation ended without verified ownership cleanup.""" + + def __init__( + self, + message: str, + process: subprocess.Popen[bytes], + original_error: BaseException | None, + cleanup_error: BaseException, + ) -> None: + super().__init__(message) + self.process = process + self.original_error = original_error + self.cleanup_error = cleanup_error + + +def _communicate_bounded( + process: subprocess.Popen[bytes], + input_bytes: bytes, + *, + deadline: float, + max_output_bytes: int, +) -> tuple[bytes, bytes]: + """Exchange bytes with a child without allowing unbounded pipe buffering.""" + if max_output_bytes < 1: + raise ValueError("header parser output limit must be positive") + if process.stdin is None or process.stdout is None or process.stderr is None: + raise RuntimeError("header parser pipes are unavailable") + + stdin_fd = process.stdin.fileno() + stdout_fd = process.stdout.fileno() + stderr_fd = process.stderr.fileno() + for fd in (stdin_fd, stdout_fd, stderr_fd): + os.set_blocking(fd, False) + + pending = memoryview(input_bytes) + offset = 0 + stdin_open = True + stdout_open = True + stderr_open = True + stdout = bytearray() + stderr = bytearray() + + def close_stdin() -> None: + nonlocal stdin_open + if not stdin_open: + return + stdin_open = False + try: + process.stdin.close() + except OSError: + pass + + while stdin_open or stdout_open or stderr_open: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("timed out checking the Lean header") + readable_fds: list[int] = [] + if stdout_open: + readable_fds.append(stdout_fd) + if stderr_open: + readable_fds.append(stderr_fd) + writable_fds = [stdin_fd] if stdin_open else [] + try: + readable, writable, _ = select.select( + readable_fds, + writable_fds, + [], + remaining, + ) + except InterruptedError: + continue + if not readable and not writable: + raise TimeoutError("timed out checking the Lean header") + + for fd, target in ((stdout_fd, stdout), (stderr_fd, stderr)): + if fd not in readable: + continue + try: + chunk = os.read(fd, 65536) + except BlockingIOError: + continue + if not chunk: + if fd == stdout_fd: + stdout_open = False + else: + stderr_open = False + continue + if len(stdout) + len(stderr) + len(chunk) > max_output_bytes: + raise ValueError( + f"Lean header parser output exceeded {max_output_bytes} bytes" + ) + target.extend(chunk) + + if stdin_fd in writable: + try: + written = os.write(stdin_fd, pending[offset : offset + 65536]) + except BlockingIOError: + continue + except BrokenPipeError: + close_stdin() + else: + if written <= 0: + close_stdin() + else: + offset += written + if offset == len(pending): + close_stdin() + + close_stdin() + if process.poll() is None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("timed out checking the Lean header") + try: + process.wait(timeout=remaining) + except subprocess.TimeoutExpired: + raise TimeoutError("timed out checking the Lean header") from None + return bytes(stdout), bytes(stderr) + + +def _decode_header_modules(stdout: bytes) -> list[str]: + """Decode the strict schemas emitted by Lean's fast import parser.""" + + def reject_constant(value: str) -> None: + raise ValueError(f"nonstandard JSON constant {value!r}") + + def reject_duplicate_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + try: + payload = json.loads( + stdout.decode("utf-8"), + parse_constant=reject_constant, + object_pairs_hook=reject_duplicate_keys, + ) + if not isinstance(payload, dict) or set(payload) != {"imports"}: + raise ValueError + entries = payload["imports"] + if not isinstance(entries, list) or len(entries) != 1: + raise ValueError + entry = entries[0] + if not isinstance(entry, dict): + raise ValueError + errors = entry.get("errors") + if not isinstance(errors, list) or not all( + isinstance(error, str) and error for error in errors + ): + raise ValueError + if errors: + raise ValueError(errors[0]) + + has_result = "result" in entry + has_imports = "imports" in entry + if has_result == has_imports: + raise ValueError + if has_result: + result = entry["result"] + if not isinstance(result, dict): + raise ValueError + imports = result.get("imports") + else: + imports = entry["imports"] + if not isinstance(imports, list): + raise ValueError + + modules: list[str] = [] + for item in imports: + if not isinstance(item, dict): + raise ValueError + module = item.get("module") + if not isinstance(module, str) or not module: + raise ValueError + if module != "Init": + modules.append(module) + return modules + except UnicodeDecodeError: + raise ValueError("unrecognized output from lean --deps-json") from None + except (TypeError, KeyError, json.JSONDecodeError): + raise ValueError("unrecognized output from lean --deps-json") from None + except ValueError as error: + if str(error): + raise + raise ValueError("unrecognized output from lean --deps-json") from None + + def _lean_header_modules( command: list[str], code: str, @@ -297,6 +494,7 @@ def _lean_header_modules( cwd: str | None, env: dict[str, str], deadline: float, + max_output_bytes: int, ) -> list[str]: """Return every module the Lean header of ``code`` imports, read by Lean itself. @@ -318,30 +516,44 @@ def _lean_header_modules( env=env, start_new_session=True, ) + original_error: BaseException | None = None + stdout = b"" + stderr = b"" try: - stdout, stderr = process.communicate(code.encode(), timeout=remaining) - except subprocess.TimeoutExpired: - _kill_subprocesses(process, process.pid) - raise TimeoutError("timed out checking the Lean header") from None - except BaseException: - _kill_subprocesses(process, process.pid) - raise + stdout, stderr = _communicate_bounded( + process, + code.encode(), + deadline=deadline, + max_output_bytes=max_output_bytes, + ) + except BaseException as error: + original_error = error + try: + _kill_subprocesses( + process, + process.pid, + time.monotonic() + DEFAULT_REPL_CLEANUP_SECONDS, + ) + except BaseException as cleanup_error: + raise _HeaderProcessCleanupError( + f"Lean header parser cleanup failed: {cleanup_error}", + process, + original_error, + cleanup_error, + ) from cleanup_error + finally: + for stream in (process.stdin, process.stdout, process.stderr): + if stream is not None and not stream.closed: + try: + stream.close() + except OSError: + pass + if original_error is not None: + raise original_error.with_traceback(original_error.__traceback__) if process.returncode != 0: detail = stderr.decode(errors="replace").strip().splitlines() raise ValueError(detail[0] if detail else f"exit status {process.returncode}") - - try: - (entry,) = json.loads(stdout)["imports"] - errors = entry["errors"] - modules = [] if errors else [item["module"] for item in entry["result"]["imports"]] - except (ValueError, TypeError, KeyError): - raise ValueError("unrecognized output from lean --deps-json") from None - if errors: - raise ValueError(str(errors[0])) - if not all(isinstance(module, str) and module for module in modules): - raise ValueError("unrecognized output from lean --deps-json") - # Every file imports Init implicitly, so importing it grants nothing. - return [module for module in modules if module != "Init"] + return _decode_header_modules(stdout) def _split_imports_and_body(code: str) -> tuple[list[str], str, int]: @@ -391,10 +603,18 @@ class LeanReplConfig: allowed_imports: frozenset[str] = ALLOWED_IMPORTS warmup_imports: frozenset[str] = WARMUP_IMPORTS - repl_command: list[str] = field(default_factory=lambda: ["lake", "exe", "repl"]) + repl_command: list[str] = field( + default_factory=lambda: ["lake", "exe", "@repl/repl"] + ) # Reports the imports of the submitted header, parsed by Lean itself. header_deps_command: list[str] = field( - default_factory=lambda: ["lake", "env", "lean", "--deps-json", "/dev/stdin"] + default_factory=lambda: [ + "lake", + "env", + sys.executable, + "-c", + _LEAN_HEADER_LAUNCHER, + ] ) # stdout is capped per response. stderr has no protocol framing, so its @@ -599,7 +819,7 @@ def __init__(self, message: str, result: dict[str, Any]) -> None: class LeanRepl: """Lean REPL process manager. - Manages a ``lake exe repl`` subprocess with non-blocking I/O, + Manages a ``lake exe @repl/repl`` subprocess with non-blocking I/O, a preloaded import environment, and automatic restart on failure. """ @@ -852,25 +1072,57 @@ def remaining() -> float: self.config.validate_imports and self._allowed_import_roots is not None ): - # The whole header reaches Lean here, so let Lean parse it. + # Validate the submitted header before adding warmup imports. + # Prefixing a `module` or `prelude` header changes what Lean + # recognizes as a header and could otherwise hide imports. env = _inherit_clean_env() env.update(self.config.env) try: header_modules = _lean_header_modules( self.config.header_deps_command, - command, + code, cwd=self.cwd, env=env, deadline=deadline, + max_output_bytes=self.config.max_buffer_bytes, ) + except _HeaderProcessCleanupError as error: + self.process = error.process + self._process_group_id = error.process.pid + self._retire_pending = True + original_error = error.original_error + cleanup_error = error.cleanup_error + if original_error is not None and not isinstance( + original_error, Exception + ): + note = str(error) + add_note = getattr(original_error, "add_note", None) + if add_note is not None: + add_note(note) + raise original_error.with_traceback( + original_error.__traceback__ + ) + if not isinstance(cleanup_error, Exception): + if original_error is not None: + note = ( + f"Header validation also failed: {original_error}" + ) + add_note = getattr(cleanup_error, "add_note", None) + if add_note is not None: + add_note(note) + raise cleanup_error.with_traceback( + cleanup_error.__traceback__ + ) + raise except ValueError as error: result = {"repl_error": f"Rejected Lean header: {error}"} header_modules = [] submitted_roots = { - module.split(".")[0] for module in header_modules + module.split(".")[0] + for module in (*header_modules, *added_imports) } disallowed = submitted_roots - self._allowed_import_roots - if disallowed: + if disallowed and result is None: result = { "repl_error": ( f"Disallowed imports: {', '.join(sorted(disallowed))}. " diff --git a/skills/setup/SKILL.md b/skills/setup/SKILL.md index 1a3f69fe..4d85fdfd 100644 --- a/skills/setup/SKILL.md +++ b/skills/setup/SKILL.md @@ -42,7 +42,10 @@ bash "/scripts/make_project.sh" \ For a new or incomplete repository: - create or repair a buildable Lean project with matching `lean-toolchain` and - Mathlib revisions; and + Mathlib revisions; +- declare `leanprover-community/repl` at an immutable revision tested against + that exact Lean toolchain, then verify `lake build @repl/repl`; never guess a + compatible revision; and - write the blueprint vault, site configuration, and CI with `autoform init`. `autoform init` is the whole vault: `blueprint/` with its landing page, @@ -91,6 +94,7 @@ then run the publication sequence: ```bash lake exe cache get # skip only when the project has no Mathlib dependency lake build +lake build @repl/repl ``` Then validate, visualize, render, and strict-build the site, keeping diff --git a/skills/setup/assets/cabannes-thesis-project/README.md b/skills/setup/assets/cabannes-thesis-project/README.md index 2d9dfed7..9a7b1c58 100644 --- a/skills/setup/assets/cabannes-thesis-project/README.md +++ b/skills/setup/assets/cabannes-thesis-project/README.md @@ -9,7 +9,8 @@ slice as a handoff example. Developed with [AutoformBot](https://github.com/facebookresearch/autoform-bot). - `lean-toolchain`, `lakefile.toml`, and `CabannesThesis/` pin matching stable - Lean and Mathlib `v4.32.2` releases. + Lean and Mathlib `v4.32.2` releases, plus a tested immutable revision of the + upstream Lean REPL used by Autoform's proof tools. - `blueprint/` is an Obsidian-compatible Markdown vault with roadmap, coverage, sources, and a seven-node theorem DAG spanning two formalization chapters. - `mkdocs.yml` builds the `autoform render` output as a leanblueprint-styled diff --git a/skills/setup/assets/cabannes-thesis-project/lake-manifest.json b/skills/setup/assets/cabannes-thesis-project/lake-manifest.json new file mode 100644 index 00000000..b4799cf9 --- /dev/null +++ b/skills/setup/assets/cabannes-thesis-project/lake-manifest.json @@ -0,0 +1,106 @@ +{"version": "1.2.0", + "packagesDir": ".lake/packages", + "packages": + [{"url": "https://github.com/leanprover-community/repl.git", + "type": "git", + "subDir": null, + "scope": "", + "rev": "68a3b3a059787a7db44fb1e6281e4a657efee470", + "name": "repl", + "manifestFile": "lake-manifest.json", + "inputRev": "68a3b3a059787a7db44fb1e6281e4a657efee470", + "inherited": false, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover-community/mathlib4.git", + "type": "git", + "subDir": null, + "scope": "", + "rev": "905b95818eb32af7874a58b427f50c1711a5e96c", + "name": "mathlib", + "manifestFile": "lake-manifest.json", + "inputRev": "v4.32.2", + "inherited": false, + "configFile": "lakefile.lean"}, + {"url": "https://github.com/leanprover-community/plausible", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "e12c1910fe855cbfc38803cd4e55543906d5fa62", + "name": "plausible", + "manifestFile": "lake-manifest.json", + "inputRev": "main", + "inherited": true, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover-community/LeanSearchClient", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "c5d5b8fe6e5158def25cd28eb94e4141ad97c843", + "name": "LeanSearchClient", + "manifestFile": "lake-manifest.json", + "inputRev": "main", + "inherited": true, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover-community/import-graph", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "7e9612bf0b9ee66db3cb5b9988a35afc706f5a12", + "name": "importGraph", + "manifestFile": "lake-manifest.json", + "inputRev": "main", + "inherited": true, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover-community/ProofWidgets4", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "6e311e2a844da9b2cc3971187df2fe0066947b93", + "name": "proofwidgets", + "manifestFile": "lake-manifest.json", + "inputRev": "main", + "inherited": true, + "configFile": "lakefile.lean"}, + {"url": "https://github.com/leanprover-community/aesop", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "a7dbf0c63b694e47f425f3dcddbc0e178bb432d3", + "name": "aesop", + "manifestFile": "lake-manifest.json", + "inputRev": "master", + "inherited": true, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover-community/quote4", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "38d591e778f100aec9762bb582f9c7f55f50e9dc", + "name": "Qq", + "manifestFile": "lake-manifest.json", + "inputRev": "master", + "inherited": true, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover-community/batteries", + "type": "git", + "subDir": null, + "scope": "leanprover-community", + "rev": "023ce7d62a0531e22a5331e20b587817a80d49ff", + "name": "batteries", + "manifestFile": "lake-manifest.json", + "inputRev": "main", + "inherited": true, + "configFile": "lakefile.toml"}, + {"url": "https://github.com/leanprover/lean4-cli", + "type": "git", + "subDir": null, + "scope": "leanprover", + "rev": "88679d088c9720c27ebdf2ba4dafe17341747f94", + "name": "Cli", + "manifestFile": "lake-manifest.json", + "inputRev": "v4.32.0", + "inherited": true, + "configFile": "lakefile.toml"}], + "name": "CabannesThesis", + "lakeDir": ".lake", + "fixedToolchain": false} diff --git a/skills/setup/assets/cabannes-thesis-project/lakefile.toml b/skills/setup/assets/cabannes-thesis-project/lakefile.toml index bc5fe5f2..1d005297 100644 --- a/skills/setup/assets/cabannes-thesis-project/lakefile.toml +++ b/skills/setup/assets/cabannes-thesis-project/lakefile.toml @@ -7,6 +7,11 @@ name = "mathlib" git = "https://github.com/leanprover-community/mathlib4.git" rev = "v4.32.2" +[[require]] +name = "repl" +git = "https://github.com/leanprover-community/repl.git" +rev = "68a3b3a059787a7db44fb1e6281e4a657efee470" + [[lean_lib]] name = "CabannesThesis" srcDir = "src" diff --git a/tests/fixtures/repl-smoke/Mathlib.lean b/tests/fixtures/repl-smoke/Mathlib.lean new file mode 100644 index 00000000..f335d1c6 --- /dev/null +++ b/tests/fixtures/repl-smoke/Mathlib.lean @@ -0,0 +1 @@ +/- A lightweight stand-in for the runtime's default warmup import. -/ diff --git a/tests/fixtures/repl-smoke/lakefile.toml b/tests/fixtures/repl-smoke/lakefile.toml index 6e845fe5..037657c9 100644 --- a/tests/fixtures/repl-smoke/lakefile.toml +++ b/tests/fixtures/repl-smoke/lakefile.toml @@ -5,3 +5,6 @@ version = "0.1.0" name = "repl" git = "https://github.com/leanprover-community/repl.git" rev = "68a3b3a059787a7db44fb1e6281e4a657efee470" + +[[lean_lib]] +name = "Mathlib" diff --git a/tests/fixtures/repl-smoke/lean-toolchain b/tests/fixtures/repl-smoke/lean-toolchain index 94b9f495..0ec5999c 100644 --- a/tests/fixtures/repl-smoke/lean-toolchain +++ b/tests/fixtures/repl-smoke/lean-toolchain @@ -1 +1 @@ -leanprover/lean4:v4.32.0 +leanprover/lean4:v4.32.2 diff --git a/tests/test_real_repl.py b/tests/test_real_repl.py index f6bb2d78..fe0157b8 100644 --- a/tests/test_real_repl.py +++ b/tests/test_real_repl.py @@ -3,12 +3,13 @@ from __future__ import annotations import os +from contextlib import suppress from pathlib import Path import pytest +from servers.lean_client import LeanRuntimeClient, LeanRuntimeUnavailable from servers.repl.core import LeanRepl, LeanReplConfig -from servers.repl.pool import LeanReplPool, LeanReplPoolConfig REPL_FIXTURE = Path(__file__).parent / "fixtures" / "repl-smoke" @@ -22,7 +23,7 @@ def test_disposable_call_matches_the_pinned_repl_protocol(): repl = LeanRepl( LeanReplConfig( cwd=str(REPL_FIXTURE), - repl_command=["lake", "exe", "repl"], + repl_command=["lake", "exe", "@repl/repl"], warmup_imports=frozenset(), validate_imports=False, ) @@ -49,18 +50,19 @@ def test_disposable_call_matches_the_pinned_repl_protocol(): ((), "/- note -/\nimport Init.Data\n#check Nat", "Disallowed imports: Init"), ((), "import REPL import Init.Data\n#check Nat", "Disallowed imports: Init"), ((), "module\npublic import Init.Data\n", "Disallowed imports: Init"), + (("REPL",), "module\npublic import Init.Data\n", "Disallowed imports: Init"), (("REPL",), "/- note -/ import Init.Data\n#check Nat", "Disallowed imports: Init"), ((), "import NotAllowlisted.Mod\n", "Disallowed imports: NotAllowlisted"), ((), "import «REPL.X»\n", "Disallowed imports: «REPL"), ((), "import «REPL\n", "Rejected Lean header"), - ((), "import REPL\ntheorem autoform_header_probe : True := trivial", None), + ((), "import REPL.Frontend\n#check Nat", None), ], ) def test_disposable_imports_are_checked_by_lean_itself(warmup, code, expected_error): repl = LeanRepl( LeanReplConfig( cwd=str(REPL_FIXTURE), - repl_command=["lake", "exe", "repl"], + repl_command=["lake", "exe", "@repl/repl"], allowed_imports=frozenset({"REPL"}), warmup_imports=frozenset(warmup), ) @@ -70,6 +72,10 @@ def test_disposable_imports_are_checked_by_lean_itself(warmup, code, expected_er if expected_error is None: assert "repl_error" not in response + assert not any( + message["severity"] == "error" + for message in response.get("messages", []) + ) else: assert expected_error in response["repl_error"] assert repl.is_clean() @@ -79,22 +85,30 @@ def test_disposable_imports_are_checked_by_lean_itself(warmup, code, expected_er os.environ.get("AUTOFORM_RUN_REAL_REPL_TESTS") != "1", reason="set AUTOFORM_RUN_REAL_REPL_TESTS=1 to run the pinned REPL integration", ) -def test_pool_calls_do_not_share_lean_state(): - pool = LeanReplPool( - LeanReplPoolConfig( - cwd=str(REPL_FIXTURE), - repl_command=["lake", "exe", "repl"], - allowed_imports=frozenset({"REPL"}), - warmup_imports=frozenset(), - num_repls=1, - ) +def test_runtime_calls_do_not_share_lean_state(runtime_dir, monkeypatch): + monkeypatch.setenv("AUTOFORM_REPL_TOTAL_WORKERS", "1") + monkeypatch.setenv("AUTOFORM_REPL_WORKERS_PER_PROJECT", "1") + client = LeanRuntimeClient( + socket_path=runtime_dir / "real-repl.sock", + response_timeout=300, + startup_timeout=30, ) - declaration = "theorem autoform_isolation_probe : True := trivial" + declaration = "theorem autoform_isolation_probe : True := True.intro" try: - responses = [pool.run(declaration, timeout=180) for _ in range(2)] + responses = [ + client.request( + "repl.run", + { + "project_dir": str(REPL_FIXTURE), + "code": declaration, + "timeout": 180, + }, + ) + for _ in range(2) + ] finally: - pool.shutdown() + with suppress(LeanRuntimeUnavailable): + client.stop() for response in responses: - assert "repl_error" not in response - assert not any(m["severity"] == "error" for m in response.get("messages", [])) + assert response == "Compiles successfully" diff --git a/tests/test_repl_core_protocol.py b/tests/test_repl_core_protocol.py index a0427474..c6588c69 100644 --- a/tests/test_repl_core_protocol.py +++ b/tests/test_repl_core_protocol.py @@ -909,7 +909,7 @@ def close(*, deadline=None): def _fake_header_deps(stdout: str, returncode: int = 0, stderr: str = "") -> list[str]: - """Stand in for ``lake env lean --deps-json /dev/stdin`` with fixed output.""" + """Stand in for the toolchain's ``lean --deps-json`` with fixed output.""" script = ( "import sys; sys.stdin.read(); " f"sys.stdout.write({stdout!r}); sys.stderr.write({stderr!r}); " @@ -928,13 +928,19 @@ def _deps_json(*modules: str, errors: tuple[str, ...] = ()) -> str: return json.dumps({"imports": [entry]}) -def _header_modules(command: list[str], deadline: float | None = None) -> list[str]: +def _header_modules( + command: list[str], + deadline: float | None = None, + *, + max_output_bytes: int = 1024 * 1024, +) -> list[str]: return repl_core._lean_header_modules( command, "import Mathlib", cwd=None, env=dict(os.environ), deadline=time.monotonic() + 10 if deadline is None else deadline, + max_output_bytes=max_output_bytes, ) @@ -946,6 +952,64 @@ def test_header_check_returns_the_modules_lean_reports(): assert _header_modules(command) == ["Mathlib.Tactic", "«Mathlib.X»", "Init.Data"] +def test_header_check_accepts_the_legacy_lean_schema(): + output = json.dumps( + { + "imports": [ + { + "errors": [], + "imports": [ + {"module": "Init", "importAll": False}, + {"module": "Mathlib.Tactic", "importAll": False}, + ], + } + ] + } + ) + + assert _header_modules(_fake_header_deps(output)) == ["Mathlib.Tactic"] + + +def test_default_header_parser_selects_lean_from_the_lake_toolchain(): + command = repl_core.LeanReplConfig().header_deps_command + + assert command[:3] == ["lake", "env", sys.executable] + assert "LEAN_SYSROOT" in command[-1] + + +def test_header_parser_launcher_ignores_a_path_shadow(tmp_path): + trusted_bin = tmp_path / "toolchain" / "bin" + shadow_bin = tmp_path / "shadow" + trusted_bin.mkdir(parents=True) + shadow_bin.mkdir() + + def write_parser(path, module): + path.write_text( + f"#!{sys.executable}\n" + "import sys\n" + "sys.stdin.read()\n" + f"sys.stdout.write({_deps_json(module)!r})\n" + ) + path.chmod(0o755) + + write_parser(trusted_bin / "lean", "Mathlib") + write_parser(shadow_bin / "lean", "Unsafe") + env = dict(os.environ) + env["LEAN_SYSROOT"] = str(trusted_bin.parent) + env["PATH"] = f"{shadow_bin}{os.pathsep}{env.get('PATH', '')}" + + modules = repl_core._lean_header_modules( + [sys.executable, "-c", repl_core._LEAN_HEADER_LAUNCHER], + "import Mathlib", + cwd=None, + env=env, + deadline=time.monotonic() + 10, + max_output_bytes=1024 * 1024, + ) + + assert modules == ["Mathlib"] + + @pytest.mark.parametrize( ("command", "message"), [ @@ -958,6 +1022,20 @@ def test_header_check_returns_the_modules_lean_reports(): (_fake_header_deps("not json"), "unrecognized output"), (_fake_header_deps('{"imports": []}'), "unrecognized output"), (_fake_header_deps('{"imports": [{"errors": [], "result": {}}]}'), "unrecognized output"), + ( + _fake_header_deps( + '{"imports": [{"errors": null, "result": {"imports": []}}]}' + ), + "unrecognized output", + ), + ( + _fake_header_deps( + '{"imports": [{"errors": [], "result": {"imports": {}}}]}' + ), + "unrecognized output", + ), + (_fake_header_deps('{"imports": [], "imports": []}'), "duplicate JSON key"), + (_fake_header_deps('{"imports": NaN}'), "nonstandard JSON constant"), (_fake_header_deps(_deps_json("")), "unrecognized output"), ], ) @@ -976,6 +1054,130 @@ def test_header_check_kills_a_command_that_outlives_the_deadline(): assert time.monotonic() - started < 5 +def test_header_check_rejects_output_over_the_configured_limit(): + command = _fake_header_deps("x" * 1025) + + with pytest.raises(ValueError, match="output exceeded 1024 bytes"): + _header_modules(command, max_output_bytes=1024) + + +def test_header_check_verifies_descendant_cleanup_after_success(tmp_path): + process_record = tmp_path / "processes.json" + script = ( + "import json, os, subprocess, sys, time; " + "child = subprocess.Popen([sys.executable, '-c', " + "'import time; time.sleep(30)'], stdout=subprocess.DEVNULL, " + "stderr=subprocess.DEVNULL); " + f"open({str(process_record)!r}, 'w').write(json.dumps(" + "{'child': child.pid, 'group': os.getpgrp()})); " + "sys.stdin.read(); " + f"sys.stdout.write({_deps_json('Mathlib')!r})" + ) + + assert _header_modules([sys.executable, "-c", script]) == ["Mathlib"] + process_ids = json.loads(process_record.read_text()) + assert not repl_core._process_group_has_live_members(process_ids["group"]) + + +def test_run_disposable_adopts_header_process_when_cleanup_must_be_retried( + monkeypatch, +): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + warmup_imports=frozenset(), + header_deps_command=_fake_header_deps(_deps_json("Mathlib")), + ) + ) + real_kill_subprocesses = repl_core._kill_subprocesses + cleanup_calls = 0 + + def fail_once(process, process_group_id, deadline=None): + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + raise RuntimeError("injected cleanup failure") + return real_kill_subprocesses(process, process_group_id, deadline) + + monkeypatch.setattr(repl_core, "_kill_subprocesses", fail_once) + monkeypatch.setattr( + repl, + "start", + lambda *args, **kwargs: pytest.fail("cleanup failure must not start Lean"), + ) + + response = repl.run_disposable("import Mathlib\n#check Nat") + + assert "header parser cleanup failed" in response["repl_error"] + assert cleanup_calls == 2 + assert repl.is_clean() + + +def test_run_disposable_preserves_header_cleanup_cancellation(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + warmup_imports=frozenset(), + header_deps_command=_fake_header_deps(_deps_json("Mathlib")), + ) + ) + real_kill_subprocesses = repl_core._kill_subprocesses + cleanup_calls = 0 + + def cancel_once(process, process_group_id, deadline=None): + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + raise asyncio.CancelledError("cancel header cleanup") + return real_kill_subprocesses(process, process_group_id, deadline) + + monkeypatch.setattr(repl_core, "_kill_subprocesses", cancel_once) + + with pytest.raises(asyncio.CancelledError, match="cancel header cleanup"): + repl.run_disposable("import Mathlib\n#check Nat") + + assert cleanup_calls == 2 + assert repl.is_clean() + + +def test_run_disposable_preserves_header_request_cancellation_when_cleanup_cancels( + monkeypatch, +): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + warmup_imports=frozenset(), + header_deps_command=[sys.executable, "-c", "import time; time.sleep(30)"], + ) + ) + real_kill_subprocesses = repl_core._kill_subprocesses + cleanup_calls = 0 + + monkeypatch.setattr( + repl_core, + "_communicate_bounded", + lambda *args, **kwargs: (_ for _ in ()).throw( + KeyboardInterrupt("cancel header request") + ), + ) + + def cancel_once(process, process_group_id, deadline=None): + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + raise asyncio.CancelledError("cancel header cleanup") + return real_kill_subprocesses(process, process_group_id, deadline) + + monkeypatch.setattr(repl_core, "_kill_subprocesses", cancel_once) + + with pytest.raises(KeyboardInterrupt, match="cancel header request") as raised: + repl.run_disposable("import Mathlib\n#check Nat") + + if hasattr(raised.value, "add_note"): + assert raised.value.__notes__ == [ + "Lean header parser cleanup failed: cancel header cleanup" + ] + assert cleanup_calls == 2 + assert repl.is_clean() + + @pytest.mark.parametrize( ("deps_output", "expected_error"), [ @@ -1006,7 +1208,29 @@ def test_run_disposable_rejects_what_lean_reports_before_starting( assert repl.process is None -def test_run_disposable_checks_the_header_after_adding_warmup_imports(monkeypatch): +def test_run_disposable_checks_warmup_import_roots(monkeypatch): + repl = repl_core.LeanRepl( + repl_core.LeanReplConfig( + allowed_imports=frozenset({"Mathlib"}), + warmup_imports=frozenset({"Mathlib.Tactic"}), + header_deps_command=_fake_header_deps(_deps_json()), + ) + ) + monkeypatch.setattr(repl, "start", lambda *args, **kwargs: None) + monkeypatch.setattr( + repl, + "_run", + lambda *args, **kwargs: {"env": 1, "messages": [], "sorries": []}, + ) + + response = repl.run_disposable("#check Nat") + + assert "repl_error" not in response + + +def test_run_disposable_checks_the_submitted_header_before_adding_warmup_imports( + monkeypatch, +): repl = repl_core.LeanRepl( repl_core.LeanReplConfig( allowed_imports=frozenset({"Mathlib"}), @@ -1023,7 +1247,7 @@ def header_modules(command, code, **kwargs): repl.run_disposable("/- note -/ import Unsafe\n#check Nat") - assert checked == ["import Mathlib\n/- note -/ import Unsafe\n#check Nat"] + assert checked == ["/- note -/ import Unsafe\n#check Nat"] class _PipeProcess: diff --git a/tests/test_shared_lean_runtime.py b/tests/test_shared_lean_runtime.py index b1b10774..162837fe 100644 --- a/tests/test_shared_lean_runtime.py +++ b/tests/test_shared_lean_runtime.py @@ -47,7 +47,7 @@ def runtime_config(**overrides): "total_repl_workers": 2, "repl_workers_per_project": 1, "repl_project_limit": 2, - "repl_command": ("lake", "exe", "repl"), + "repl_command": ("lake", "exe", "@repl/repl"), "lsp_command": ("lake", "serve"), "lsp_timeout": 60.0, "max_lsp_request_seconds": 600.0, @@ -1951,6 +1951,16 @@ def test_runtime_response_rejects_invalid_utf8_as_unknown_outcome( assert isinstance(caught.value.__cause__, LeanRuntimeProtocolError) +def test_default_repl_command_uses_the_declared_dependency(monkeypatch): + monkeypatch.delenv("LEAN_REPL_CMD", raising=False) + + assert LeanRuntimeConfig.from_environment().repl_command == ( + "lake", + "exe", + "@repl/repl", + ) + + @pytest.mark.parametrize( ("name", "value", "match"), [ diff --git a/tests/test_skill_examples.py b/tests/test_skill_examples.py index 04ff8735..586d21ce 100644 --- a/tests/test_skill_examples.py +++ b/tests/test_skill_examples.py @@ -117,10 +117,40 @@ def test_setup_asset_is_a_repo_shaped_thesis_vault(repo_root: Path) -> None: assert (example / "src/CabannesThesis.lean").is_file() assert (example / "src/CabannesThesis/Basic.lean").is_file() toolchain = (example / "lean-toolchain").read_text(encoding="utf-8").strip() - manifest = tomllib.loads((example / "lakefile.toml").read_text(encoding="utf-8")) + lakefile = tomllib.loads((example / "lakefile.toml").read_text(encoding="utf-8")) assert toolchain == "leanprover/lean4:v4.32.2" - assert manifest["require"][0]["rev"] == "v4.32.2" - assert manifest["lean_lib"][0]["srcDir"] == "src" + requirements = { + requirement["name"]: requirement for requirement in lakefile["require"] + } + assert requirements["mathlib"]["rev"] == "v4.32.2" + assert requirements["repl"] == { + "name": "repl", + "git": "https://github.com/leanprover-community/repl.git", + "rev": "68a3b3a059787a7db44fb1e6281e4a657efee470", + } + assert lakefile["lean_lib"][0]["srcDir"] == "src" + manifest = json.loads( + (example / "lake-manifest.json").read_text(encoding="utf-8") + ) + repl_package = next( + package for package in manifest["packages"] if package["name"] == "repl" + ) + assert repl_package["rev"] == requirements["repl"]["rev"] + + fixture = repo_root / "tests/fixtures/repl-smoke" + assert ( + (fixture / "lean-toolchain").read_text(encoding="utf-8").strip() + == toolchain + ) + fixture_lakefile = tomllib.loads( + (fixture / "lakefile.toml").read_text(encoding="utf-8") + ) + fixture_repl = next( + requirement + for requirement in fixture_lakefile["require"] + if requirement["name"] == "repl" + ) + assert fixture_repl["rev"] == requirements["repl"]["rev"] def test_setup_asset_static_site_contract(repo_root: Path, tmp_path: Path) -> None: