Move the LSP token scan into servers/_process.py and drop the underscore from its shared names

servers/_process.py

@@ -1,4 +1,4 @@
-"""Give a Lean child a clean environment, and stop and reap its dedicated POSIX process group."""
+"""Give a Lean child a clean environment, find the processes carrying a token the caller set, and stop and reap its dedicated POSIX process group."""
 
 from __future__ import annotations
 
@@ -7,6 +7,7 @@
 import signal
 import subprocess
 import time
+from typing import Any
 
 REPL_ABORT_TERM_SECONDS = 0.5
 REPL_ABORT_KILL_SECONDS = 1.0
@@ -28,7 +29,7 @@ def _process_group_exists(process_group_id: int) -> bool:
     return True
 
 
-def _process_group_has_live_members(process_group_id: int) -> bool:
+def process_group_has_live_members(process_group_id: int) -> bool:
     """Treat an all-zombie process group as terminated after SIGKILL."""
     if not _process_group_exists(process_group_id):
         return False
@@ -49,7 +50,7 @@ def _process_group_has_live_members(process_group_id: int) -> bool:
 
 
 def _wait_for_live_process_group_exit(process_group_id: int, deadline: float) -> bool:
-    while _process_group_has_live_members(process_group_id):
+    while process_group_has_live_members(process_group_id):
         remaining = deadline - time.monotonic()
         if remaining <= 0:
             return False
@@ -77,7 +78,7 @@ def _kill_subprocesses(
 ) -> None:
     """Terminate and verify one dedicated POSIX process group."""
     if getattr(process, "returncode", None) is not None:
-        if _process_group_has_live_members(process_group_id):
+        if process_group_has_live_members(process_group_id):
             raise RuntimeError(
                 "refusing to signal a process group after its leader was reaped"
             )
@@ -98,7 +99,7 @@ def _terminate_process_group(
         term_deadline = min(term_deadline, deadline)
         kill_deadline = min(kill_deadline, deadline)
 
-    if _process_group_has_live_members(process_group_id):
+    if process_group_has_live_members(process_group_id):
         try:
             os.killpg(process_group_id, signal.SIGTERM)
         except ProcessLookupError:
@@ -107,7 +108,7 @@ def _terminate_process_group(
             raise RuntimeError("failed to terminate the Lean process group") from error
 
     _wait_for_live_process_group_exit(process_group_id, term_deadline)
-    if _process_group_has_live_members(process_group_id):
+    if process_group_has_live_members(process_group_id):
         try:
             os.killpg(process_group_id, signal.SIGKILL)
         except ProcessLookupError:
@@ -143,8 +144,44 @@ def _reap_process(
         raise RuntimeError("timed out reaping the Lean process")
 
 
-def _inherit_clean_env() -> dict[str, str]:
+def inherit_clean_env() -> dict[str, str]:
     """Return a copy of the current environment without PYTHONPATH noise."""
     env = os.environ.copy()
     env.pop("PYTHONPATH", None)
     return env
+
+
+def processes_with_token(name: str, token: str) -> list[Any]:
+    """Return the processes whose environment sets ``name`` to ``token``.
+
+    A process whose environment cannot be read is left out, so it is never
+    signalled: nothing shows that it carries the token. Lean and Lake run as
+    the user, who can read their environments; setuid programs the user runs,
+    such as ``login``, are unreadable at all times, so counting them would
+    fail every close. A process table that cannot be read fails the scan.
+    """
+    if not token:
+        # A missing token would match every process without the variable.
+        raise ValueError("a process token must be a nonempty string")
+    import psutil
+
+    try:
+        pids = psutil.pids()
+    except (OSError, psutil.Error) as error:
+        raise RuntimeError(f"cannot read the process table to find the processes carrying {name}") from error
+    found = []
+    for pid in pids:
+        if pid == os.getpid():
+            continue
+        try:
+            # The handle records the creation time before the environment is
+            # read, so it can only name the process that was read or an
+            # earlier one, which psutil then refuses to signal.
+            candidate = psutil.Process(pid)
+            environment = candidate.environ()
+        except (psutil.NoSuchProcess, psutil.AccessDenied):
+            # Exited or unreadable; a zombie reads as exited or empty.
+            continue
+        if environment.get(name) == token:
+            found.append(candidate)
+    return found

servers/batch_diagnostics.py

@@ -21,9 +21,9 @@
 from typing import Any
 
 from servers._process import (
-    _inherit_clean_env,
     _kill_subprocesses,
     _wait_for_live_process_group_exit,
+    inherit_clean_env,
 )
 
 logger = logging.getLogger(__name__)
@@ -118,7 +118,7 @@ def _run(argv: list[str], root: Path, timeout: float) -> tuple[int, bytes | None
         process = subprocess.Popen(
             argv,
             cwd=root,
-            env={**_inherit_clean_env(), "TMPDIR": scratch},
+            env={**inherit_clean_env(), "TMPDIR": scratch},
             stdin=subprocess.DEVNULL,
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,

servers/lsp/server.py

@@ -44,11 +44,12 @@
 from servers._process import (
     REPL_ABORT_KILL_SECONDS,
     REPL_ABORT_TERM_SECONDS,
-    _inherit_clean_env,
-    _process_group_has_live_members,
     _reap_process,
     _terminate_process_group,
     _wait_for_live_process_group_exit,
+    inherit_clean_env,
+    process_group_has_live_members,
+    processes_with_token,
 )
 from servers.lean_client import LeanRuntimeClient
 
@@ -120,42 +121,6 @@ class LspConfig:
     timeout: float = DEFAULT_LSP_TIMEOUT
 
 
-def _processes_with_token(token: str) -> list[Any]:
-    """Return the processes whose environment carries a session's token.
-
-    A process whose environment cannot be read is left out, so it is never
-    signalled: nothing shows that it is the session's. Lean and Lake run as
-    the user, who can read their environments; setuid programs the user runs,
-    such as ``login``, are unreadable at all times, so counting them would
-    fail every close. A process table that cannot be read fails the scan.
-    """
-    if not token:
-        # A missing token would match every process without the variable.
-        raise ValueError("a Lean LSP session token must be a nonempty string")
-    import psutil
-
-    try:
-        pids = psutil.pids()
-    except (OSError, psutil.Error) as error:
-        raise RuntimeError("cannot read the process table to find the Lean LSP session's processes") from error
-    found = []
-    for pid in pids:
-        if pid == os.getpid():
-            continue
-        try:
-            # The handle records the creation time before the environment is
-            # read, so it can only name the process that was read or an
-            # earlier one, which psutil then refuses to signal.
-            candidate = psutil.Process(pid)
-            environment = candidate.environ()
-        except (psutil.NoSuchProcess, psutil.AccessDenied):
-            # Exited or unreadable; a zombie reads as exited or empty.
-            continue
-        if environment.get(_SESSION_TOKEN_ENV) == token:
-            found.append(candidate)
-    return found
-
-
 def _running(process: Any) -> bool:
     """Whether a psutil process is still running; one that cannot be read counts as running."""
     import psutil
@@ -197,7 +162,7 @@ def start(self) -> None:
         """Start the language server process."""
         if os.name != "posix":
             raise RuntimeError("Lean LSP transport requires a POSIX platform")
-        env = _inherit_clean_env()
+        env = inherit_clean_env()
         self._token = secrets.token_hex(16)
         env[_SESSION_TOKEN_ENV] = self._token
         process: subprocess.Popen | None = None
@@ -255,7 +220,7 @@ def _retire_once(self) -> None:
                 process_group_id = process.pid
                 self._process_group_id = process_group_id
             if getattr(process, "returncode", None) is not None:
-                if _process_group_has_live_members(process_group_id):
+                if process_group_has_live_members(process_group_id):
                     raise RuntimeError(
                         "refusing to signal an LSP process group after its leader was reaped"
                     )
@@ -309,7 +274,7 @@ def _stop_processes_with_token(self) -> None:
             (signal.SIGTERM, REPL_ABORT_TERM_SECONDS),
             (signal.SIGKILL, REPL_ABORT_KILL_SECONDS),
         ):
-            processes = _processes_with_token(token)
+            processes = processes_with_token(_SESSION_TOKEN_ENV, token)
             if not processes:
                 return
             for process in processes:
@@ -320,7 +285,7 @@ def _stop_processes_with_token(self) -> None:
                 except (OSError, psutil.Error):
                     pass
             _wait_for_exit(processes, time.monotonic() + seconds)
-        if _processes_with_token(token):
+        if processes_with_token(_SESSION_TOKEN_ENV, token):
             raise RuntimeError("processes started by the Lean LSP server are still running")
 
     def _abort_process(self) -> None:
@@ -539,7 +504,7 @@ def _left_running(self, deadline: float) -> bool:
         if token is None:
             return False
         try:
-            carriers = _processes_with_token(token)
+            carriers = processes_with_token(_SESSION_TOKEN_ENV, token)
         except RuntimeError:
             return True
         outside = []

servers/repl/core.py

@@ -25,7 +25,7 @@
 from logging import getLogger
 from typing import Any, Callable
 
-from servers._process import _inherit_clean_env, _kill_subprocesses
+from servers._process import _kill_subprocesses, inherit_clean_env
 
 logger = getLogger(__name__)
 
@@ -737,7 +737,7 @@ def _spawn(self, command: list[str]) -> subprocess.Popen:
         The child leads its own session, and close() retires that whole
         process group through the id published here.
         """
-        env = _inherit_clean_env()
+        env = inherit_clean_env()
         env.update(self.config.env)
         self.process = subprocess.Popen(
             command,

tests/test_lsp_protocol.py

@@ -18,6 +18,7 @@
 import psutil
 import pytest
 
+from servers._process import processes_with_token
 from servers.lsp import server as lsp
 from tests.lean_fixture import _real_lean_project, _requires_real_lean
 
@@ -336,17 +337,19 @@ def test_abort_retries_the_token_sweep_before_reaping_the_leader(monkeypatch):
     # Each scan of the first attempt finds a process carrying the token, so
     # one is left after SIGKILL; the retry's first scan finds none.
     scans = iter([[_FakeTokenCarrier(events)]] * 3)
+    # Asserting inside the scan would only be retried by abort, so record.
+    calls: list[tuple[str, str]] = []
 
-    def scan(token):
-        assert token == "session-token"
+    def scan(name, token):
+        calls.append((name, token))
         events.append("scan")
         return next(scans, [])
 
     def reap(owned):
         events.append("reap")
         owned.wait()
 
-    monkeypatch.setattr(lsp, "_processes_with_token", scan)
+    monkeypatch.setattr(lsp, "processes_with_token", scan)
     monkeypatch.setattr(lsp, "_terminate_process_group", lambda process_group_id: events.append("group"))
     monkeypatch.setattr(lsp, "_reap_process", reap)
     monkeypatch.setattr(lsp.time, "sleep", lambda delay: None)
@@ -357,6 +360,7 @@ def reap(owned):
     # the retired group again, and the leader was reaped only after a sweep
     # found nothing left.
     assert events == ["group", "scan", "SIGTERM", "scan", "SIGKILL", "scan", "scan", "reap"]
+    assert calls == [(lsp._SESSION_TOKEN_ENV, "session-token")] * 4
     assert session.process is None
 
 
@@ -366,7 +370,7 @@ def test_token_sweep_signals_only_what_its_latest_scan_found(monkeypatch):
     session = _owned_session(_FakeProcess(), "session-token")
     signals: list[str] = []
     scans = iter([[_FakeTokenCarrier(signals)], []])
-    monkeypatch.setattr(lsp, "_processes_with_token", lambda token: next(scans, []))
+    monkeypatch.setattr(lsp, "processes_with_token", lambda name, token: next(scans, []))
     monkeypatch.setattr(lsp, "_terminate_process_group", lambda process_group_id: None)
     monkeypatch.setattr(lsp, "_reap_process", lambda owned: owned.wait())
 
@@ -539,7 +543,7 @@ class StopRetry(BaseException):
 
     monkeypatch.setattr(
         lsp,
-        "_process_group_has_live_members",
+        "process_group_has_live_members",
         lambda process_group_id: True,
     )
     monkeypatch.setattr(
@@ -575,7 +579,9 @@ def test_token_scan_finds_only_the_processes_carrying_the_token():
     }
     try:
         # Popen returns once the child has exec'd, so its environment is set.
-        assert [process.pid for process in lsp._processes_with_token(token)] == [processes["this session"].pid]
+        assert [process.pid for process in processes_with_token(lsp._SESSION_TOKEN_ENV, token)] == [
+            processes["this session"].pid
+        ]
     finally:
         for process in processes.values():
             process.kill()
@@ -605,7 +611,25 @@ def environ(self):
     monkeypatch.setattr(psutil, "pids", lambda: list(candidates))
     monkeypatch.setattr(psutil, "Process", lambda pid: candidates[pid])
 
-    assert lsp._processes_with_token("session-token") == [candidates[1]]
+    assert processes_with_token(lsp._SESSION_TOKEN_ENV, "session-token") == [candidates[1]]
+
+
+def test_token_scan_reads_the_variable_it_is_given(monkeypatch):
+    class Candidate:
+        def __init__(self, environment):
+            self.environment = environment
+
+        def environ(self):
+            return self.environment
+
+    candidates = {
+        1: Candidate({"OTHER_TOKEN": "session-token"}),
+        2: Candidate({lsp._SESSION_TOKEN_ENV: "session-token"}),
+    }
+    monkeypatch.setattr(psutil, "pids", lambda: list(candidates))
+    monkeypatch.setattr(psutil, "Process", lambda pid: candidates[pid])
+
+    assert processes_with_token("OTHER_TOKEN", "session-token") == [candidates[1]]
 
 
 def test_token_scan_fails_when_the_process_table_cannot_be_read(monkeypatch):
@@ -616,7 +640,7 @@ def unreadable():
 
     # The attempt fails, so close and abort retry instead of reaping the leader.
     with pytest.raises(RuntimeError, match="cannot read the process table"):
-        lsp._processes_with_token("session-token")
+        processes_with_token(lsp._SESSION_TOKEN_ENV, "session-token")
 
 
 @pytest.mark.parametrize(
@@ -635,14 +659,14 @@ def environ(self):
     # Only exited and unreadable processes are passed over; anything else
     # fails the attempt, so close and abort retry.
     with pytest.raises(type(error)):
-        lsp._processes_with_token("session-token")
+        processes_with_token(lsp._SESSION_TOKEN_ENV, "session-token")
 
 
 @pytest.mark.parametrize("token", [None, ""], ids=["none", "empty"])
 def test_token_scan_refuses_a_missing_token(token):
     # A missing token would match every process without the variable.
     with pytest.raises(ValueError, match="nonempty"):
-        lsp._processes_with_token(token)
+        processes_with_token(lsp._SESSION_TOKEN_ENV, token)
 
 
 def test_concurrent_abort_and_close_serialize_owned_cleanup(monkeypatch):