Finalize Lean declaration sessions

Author: morluto

src/jacobian/lean_frontend/declarations.py

@@ -8,6 +8,7 @@
 import tempfile
 import threading
 import uuid
+import weakref
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Protocol
@@ -290,6 +291,16 @@ class _SessionEntry:
     session: _DeclarationQuerySession
 
 
+def _close_declaration_sessions(
+    sessions: dict[LeanEnvironment, _SessionEntry],
+    session_locks: dict[LeanEnvironment, threading.Lock],
+) -> None:
+    """Release declaration sessions retained past their backend's lifetime."""
+    for environment, entry in list(sessions.items()):
+        with session_locks[environment]:
+            entry.session.close()
+
+
 class LeanSubprocessDeclarationBackend:
     """Reuse one catalog across bounded processes for each validated environment."""
 
@@ -312,6 +323,12 @@ def __init__(
         self._session_locks = {
             environment: threading.Lock() for environment in LeanEnvironment
         }
+        self._finalizer = weakref.finalize(
+            self,
+            _close_declaration_sessions,
+            self._sessions,
+            self._session_locks,
+        )
 
     def environment_digest(self, environment: LeanEnvironment) -> str:
         try:
@@ -450,6 +467,7 @@ def close(self) -> None:
         for environment, lock in self._session_locks.items():
             with lock:
                 self._discard_session(environment)
+        self._finalizer.detach()
 
     def _start_session(
         self,

tests/component/providers/lean/test_lean_declaration_sessions.py

@@ -2,8 +2,10 @@
 
 from __future__ import annotations
 
+import gc
 import hashlib
 import threading
+import weakref
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -152,6 +154,32 @@ def test_subprocess_backend_reuses_one_pinned_session_per_environment(
     assert session.closed
 
 
+def test_subprocess_backend_finalizer_closes_sessions_after_garbage_collection(
+    tmp_path: Path,
+) -> None:
+    session = RecordingSession(responses=[{"operation": "search", "declarations": []}])
+    backend, _ = _recording_backend(tmp_path, [session])
+    backend.query(LeanEnvironment.CORE, {"operation": "search"})
+    reference = weakref.ref(backend)
+
+    del backend
+    gc.collect()
+
+    assert reference() is None
+    assert session.closed
+
+
+def test_subprocess_backend_close_detaches_finalizer(tmp_path: Path) -> None:
+    session = RecordingSession(responses=[{"operation": "search", "declarations": []}])
+    backend, _ = _recording_backend(tmp_path, [session])
+    backend.query(LeanEnvironment.CORE, {"operation": "search"})
+
+    backend.close()
+
+    assert session.closed
+    assert not backend._finalizer.alive
+
+
 def test_subprocess_backend_serializes_concurrent_session_requests(
     tmp_path: Path,
 ) -> None: