Normalize and consolidate shared LSP calls

servers/lean_runtime.py

@@ -914,21 +914,10 @@ def dispatch(self, method: str, params: dict[str, Any]) -> Any:
             project_dir = self._string_param(params, "project_dir")
             file_path = self._string_param(params, "file_path")
             root, path = resolve_lean_file(project_dir, file_path)
-            with self.lsp_projects.lease_resolved(
+            diagnostics = self._lsp_operation(
                 root,
-                acquisition_timeout=self._acquisition_timeout(self.config.lsp_timeout),
-                creation_budget=self.lsp_creation_budget,
-            ) as session:
-                assert session is not None
-                try:
-                    diagnostics = session.get_diagnostics(str(path))
-                except Exception:
-                    if not session.is_alive():
-                        session.retire()
-                        self.lsp_projects.invalidate_resolved(root, session)
-                    raise
-                if not session.is_alive():
-                    self.lsp_projects.invalidate_resolved(root, session)
+                lambda session: session.get_diagnostics(str(path)),
+            )
             return format_lsp_diagnostics(diagnostics)
         if method == "lsp.hover":
             project_dir = self._string_param(params, "project_dir")
@@ -938,24 +927,36 @@ def dispatch(self, method: str, params: dict[str, Any]) -> Any:
             if line < 0 or character < 0:
                 raise ValueError("line and character must be nonnegative")
             root, path = resolve_lean_file(project_dir, file_path)
-            with self.lsp_projects.lease_resolved(
+            result = self._lsp_operation(
                 root,
-                acquisition_timeout=self._acquisition_timeout(self.config.lsp_timeout),
-                creation_budget=self.lsp_creation_budget,
-            ) as session:
-                assert session is not None
-                try:
-                    result = session.hover(str(path), line, character)
-                except Exception:
-                    if not session.is_alive():
-                        session.retire()
-                        self.lsp_projects.invalidate_resolved(root, session)
-                    raise
-                if not session.is_alive():
-                    self.lsp_projects.invalidate_resolved(root, session)
+                lambda session: session.hover(str(path), line, character),
+            )
             return result or "No hover information at this position."
         raise ValueError(f"unknown Lean runtime method: {method}")
 
+    def _lsp_operation(
+        self,
+        root: Path,
+        operation: Callable[[LeanLspSession], T],
+    ) -> T:
+        """Route one LSP call while preserving retained cleanup ownership."""
+        with self.lsp_projects.lease_resolved(
+            root,
+            acquisition_timeout=self._acquisition_timeout(self.config.lsp_timeout),
+            creation_budget=self.lsp_creation_budget,
+        ) as session:
+            assert session is not None
+            try:
+                result = operation(session)
+            except Exception:
+                if not session.is_alive():
+                    session.retire()
+                    self.lsp_projects.invalidate_resolved(root, session)
+                raise
+            if not session.is_alive():
+                self.lsp_projects.invalidate_resolved(root, session)
+            return result
+
     def status(self, *, include_projects: bool) -> dict[str, Any]:
         result: dict[str, Any] = {
             "running": True,

servers/lsp/server.py

@@ -15,7 +15,7 @@
 from dataclasses import dataclass, field
 from logging import getLogger
 from pathlib import Path
-from typing import Any
+from typing import Any, TypeVar
 
 from fastmcp.server import FastMCP
 
@@ -32,6 +32,7 @@
 # Bound at import: tests swap subprocess.Popen for factories, and non-Popen
 # fakes skip the psutil liveness probe.
 _POPEN_TYPE = subprocess.Popen
+T = TypeVar("T")
 
 DEFAULT_LSP_TIMEOUT = 60
 MAX_LSP_HEADER_BYTES = 16 * 1024
@@ -238,8 +239,8 @@ def is_alive(self) -> bool:
             return True
         return status not in {psutil.STATUS_ZOMBIE, psutil.STATUS_DEAD}
 
-    def get_diagnostics(self, file_path: str) -> list[dict]:
-        """Open a file and collect diagnostics from the language server."""
+    def _admitted(self, operation: Callable[[float], T]) -> T:
+        """Run one serialized operation and fence a failed protocol stream."""
         deadline = time.monotonic() + self.config.timeout
         if not self._operation_lock.acquire(timeout=self.config.timeout):
             raise LspBusyError(
@@ -254,13 +255,19 @@ def get_diagnostics(self, file_path: str) -> list[dict]:
                     f"timed out after {self.config.timeout:g}s waiting for the Lean LSP session"
                 )
             try:
-                return self._get_diagnostics(file_path, timeout=remaining)
+                return operation(remaining)
             except BaseException:
                 self._retire_pending = True
                 raise
         finally:
             self._operation_lock.release()
 
+    def get_diagnostics(self, file_path: str) -> list[dict]:
+        """Open a file and collect diagnostics from the language server."""
+        return self._admitted(
+            lambda remaining: self._get_diagnostics(file_path, timeout=remaining)
+        )
+
     def _get_diagnostics(self, file_path: str, *, timeout: float | None = None) -> list[dict]:
         path = Path(file_path).resolve()
         uri = path.as_uri()
@@ -310,32 +317,15 @@ def hover(self, file_path: str, line: int, character: int) -> str | None:
         # shared protocol stream, so they must not retire a healthy session.
         path = Path(file_path).resolve()
         content = path.read_text()
-        deadline = time.monotonic() + self.config.timeout
-        if not self._operation_lock.acquire(timeout=self.config.timeout):
-            raise LspBusyError(
-                f"timed out after {self.config.timeout:g}s waiting for the Lean LSP session"
+        return self._admitted(
+            lambda remaining: self._hover(
+                path,
+                content,
+                line,
+                character,
+                timeout=remaining,
             )
-        try:
-            if self._retire_pending:
-                raise LspProtocolError("Lean LSP session is retiring after a failed operation")
-            remaining = deadline - time.monotonic()
-            if remaining <= 0:
-                raise LspBusyError(
-                    f"timed out after {self.config.timeout:g}s waiting for the Lean LSP session"
-                )
-            try:
-                return self._hover(
-                    path,
-                    content,
-                    line,
-                    character,
-                    timeout=remaining,
-                )
-            except BaseException:
-                self._retire_pending = True
-                raise
-        finally:
-            self._operation_lock.release()
+        )
 
     def _hover(
         self,
@@ -570,7 +560,7 @@ def _read_message(self, timeout: float = 5) -> dict | None:
 
         try:
             message = json.loads(bytes(body).decode("utf-8"))
-        except (UnicodeDecodeError, json.JSONDecodeError) as error:
+        except (ValueError, RecursionError) as error:
             raise LspProtocolError("LSP emitted an invalid JSON body") from error
         if not isinstance(message, dict):
             raise LspProtocolError("LSP JSON-RPC message is not an object")

tests/test_lsp_protocol.py

@@ -369,6 +369,60 @@ def notify(method, params, **kwargs):
             session.hover(str(source), 0, 0)
 
 
+@pytest.mark.parametrize("operation", ["diagnostics", "hover"])
+def test_post_admission_budget_exhaustion_keeps_session_healthy(
+    tmp_path, monkeypatch, operation
+):
+    source = tmp_path / "Test.lean"
+    source.write_text("#check Nat\n")
+    session = _owned_session(_FakeProcess())
+    session.config.timeout = 0
+    monkeypatch.setattr(lsp.time, "monotonic", lambda: 0.0)
+    monkeypatch.setattr(
+        session,
+        "_get_diagnostics",
+        lambda *args, **kwargs: pytest.fail("expired admission ran diagnostics"),
+    )
+    monkeypatch.setattr(
+        session,
+        "_hover",
+        lambda *args, **kwargs: pytest.fail("expired admission ran hover"),
+    )
+
+    with pytest.raises(lsp.LspBusyError, match="waiting for the Lean LSP session"):
+        if operation == "diagnostics":
+            session.get_diagnostics(str(source))
+        else:
+            session.hover(str(source), 0, 0)
+
+    assert session.is_alive() is True
+    assert session._retire_pending is False
+
+
+def test_document_close_budget_exhaustion_retires_session(tmp_path, monkeypatch):
+    source = tmp_path / "Test.lean"
+    source.write_text("#check Nat\n")
+    session = lsp.LeanLspSession(lsp.LspConfig(timeout=60))
+    clock = {"now": 0.0}
+    notifications = []
+    monkeypatch.setattr(lsp.time, "monotonic", lambda: clock["now"])
+    monkeypatch.setattr(
+        session,
+        "_send_notification",
+        lambda method, params, **kwargs: notifications.append(method),
+    )
+
+    def finish_at_deadline(uri, timeout):
+        clock["now"] = 60.0
+        return []
+
+    monkeypatch.setattr(session, "_collect_diagnostics", finish_at_deadline)
+
+    assert session.get_diagnostics(str(source)) == []
+    assert notifications == ["textDocument/didOpen"]
+    assert session._retire_pending is True
+
+
 def test_hover_input_decode_failure_keeps_healthy_session(tmp_path):
     source = tmp_path / "Invalid.lean"
     source.write_bytes(b"\xff")
@@ -870,6 +924,36 @@ def test_lsp_queue_wait_is_bounded_by_session_timeout(tmp_path: Path):
         session._operation_lock.release()
 
 
+def test_lsp_json_integer_limit_is_reported_as_a_protocol_error():
+    body = b'{"jsonrpc":"2.0","id":' + (b"9" * 5000) + b"}"
+    frame = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + body
+    read_fd, write_fd = os.pipe()
+    reader = os.fdopen(read_fd, "rb", buffering=0)
+    session = lsp.LeanLspSession(lsp.LspConfig())
+    session.process = SimpleNamespace(stdout=reader)
+    try:
+        os.write(write_fd, frame)
+        with pytest.raises(lsp.LspProtocolError, match="invalid JSON body"):
+            session._read_message(timeout=1)
+    finally:
+        os.close(write_fd)
+        reader.close()
+
+
+def test_lsp_json_nesting_limit_is_reported_as_a_protocol_error(tmp_path: Path):
+    # 200,000 levels overrun the JSON decoder's recursion limit (about 10,000
+    # on Python 3.13, 1,000 on 3.10) in a body far under MAX_LSP_MESSAGE_BYTES.
+    # A pipe cannot buffer the whole frame, so the session reads it from a file.
+    body = b"[" * 200_000 + b"]" * 200_000
+    frame = tmp_path / "frame"
+    frame.write_bytes(f"Content-Length: {len(body)}\r\n\r\n".encode("ascii") + body)
+    session = lsp.LeanLspSession(lsp.LspConfig())
+    with frame.open("rb", buffering=0) as reader:
+        session.process = SimpleNamespace(stdout=reader)
+        with pytest.raises(lsp.LspProtocolError, match="invalid JSON body"):
+            session._read_message(timeout=1)
+
+
 @pytest.mark.parametrize(
     "partial",
     [