Keep visualization worker alive until explicit stop

src/docforge/visualization.py

@@ -3,15 +3,20 @@
 from __future__ import annotations
 
 import atexit
+import fcntl
 import json
 import os
 import secrets
+import signal
 import socket
 import sqlite3
-import subprocess
 import sys
+import tempfile
 import threading
+import time
+import urllib.error
 import urllib.parse
+import urllib.request
 from collections.abc import Generator
 from contextlib import contextmanager
 from http import HTTPStatus
@@ -22,6 +27,7 @@ from typing import cast
 
 from .errors import DocForgeError
 from .index import APPLICATION_ID, INDEX_SCHEMA_VERSION, ProjectIndex, re_tokenize
+from .project import project_root_fingerprint
 
 VISUALIZATION_TEMPLATE = "graph-browser@8"
 DEFAULT_EDGE_LIMIT = 100
@@ -29,6 +35,9 @@ MAX_EDGE_LIMIT = 400
 DEFAULT_INITIAL_GRACE_SECONDS = 120.0
 DEFAULT_LEASE_SECONDS = 180.0
 LEASE_MONITOR_INTERVAL_SECONDS = 1.0
+VISUALIZATION_REGISTRY_NAME = ".visualization.json"
+VISUALIZATION_LOCK_NAME = ".visualization.lock"
+VISUALIZATION_RUNTIME = "persistent-worker@1"
 
 
 class _VisualizationHttpServer(ThreadingHTTPServer):
@@ -359,17 +368,23 @@ class VisualizationRunner:
         snapshot_spec: dict[str, object] | None = None,
         token: str | None = None,
         register_atexit: bool = True,
+        persistent: bool = False,
         initial_grace_seconds: float = DEFAULT_INITIAL_GRACE_SECONDS,
         lease_seconds: float = DEFAULT_LEASE_SECONDS,
         monitor_interval_seconds: float = LEASE_MONITOR_INTERVAL_SECONDS,
     ) -> None:
         if (index is None) == (snapshot_spec is None):
             raise ValueError("Provide exactly one visualization index or snapshot")
-        if initial_grace_seconds <= 0 or lease_seconds <= 0 or monitor_interval_seconds <= 0:
+        if (
+            (not persistent and initial_grace_seconds <= 0)
+            or lease_seconds <= 0
+            or monitor_interval_seconds <= 0
+        ):
             raise ValueError("Visualization lease durations must be positive")
         self.index = index
         self._snapshot_spec = snapshot_spec
         self._register_atexit = register_atexit
+        self.persistent = persistent
         self.initial_grace_seconds = initial_grace_seconds
         self.lease_seconds = lease_seconds
         self.monitor_interval_seconds = monitor_interval_seconds
@@ -455,12 +470,13 @@ class VisualizationRunner:
                     daemon=False,
                 )
                 self._thread.start()
-                self._lease_thread = threading.Thread(
-                    target=self._monitor_lease,
-                    name="docforge-visualization-lease",
-                    daemon=True,
-                )
-                self._lease_thread.start()
+                if not self.persistent:
+                    self._lease_thread = threading.Thread(
+                        target=self._monitor_lease,
+                        name="docforge-visualization-lease",
+                        daemon=True,
+                    )
+                    self._lease_thread.start()
                 if self._register_atexit and not self._atexit_registered:
                     atexit.register(self.stop)
                     self._atexit_registered = True
@@ -486,11 +502,15 @@ class VisualizationRunner:
             "template": VISUALIZATION_TEMPLATE,
             "read_only": True,
             "project_bound": True,
-            "lifetime": {
-                "policy": "browser_lease",
-                "initial_grace_seconds": self.initial_grace_seconds,
-                "lease_seconds": self.lease_seconds,
-            },
+            "lifetime": (
+                {"policy": "explicit_stop"}
+                if self.persistent
+                else {
+                    "policy": "browser_lease",
+                    "initial_grace_seconds": self.initial_grace_seconds,
+                    "lease_seconds": self.lease_seconds,
+                }
+            ),
             "target": {
                 "node_id": node_id,
                 "query": query,
@@ -730,26 +750,207 @@ class VisualizationRunner:
         )
 
 
-class DetachedVisualizationRunner:
-    """Launch a viewer worker that survives a short-lived MCP transport process."""
+class PersistentVisualizationRunner:
+    """Run one project-bound browser until an explicit DocForge stop request."""
 
-    def __init__(
-        self,
-        index: ProjectIndex,
-        *,
-        owner_pid: int | None = None,
-        initial_grace_seconds: float = DEFAULT_INITIAL_GRACE_SECONDS,
-        lease_seconds: float = DEFAULT_LEASE_SECONDS,
-        monitor_interval_seconds: float = LEASE_MONITOR_INTERVAL_SECONDS,
-    ) -> None:
-        if owner_pid is not None and owner_pid <= 1:
-            raise ValueError("Visualization owner PID must identify a live user process")
+    def __init__(self, index: ProjectIndex) -> None:
         self.index = index
-        self.owner_pid = owner_pid if owner_pid is not None else os.getppid()
-        self.initial_grace_seconds = initial_grace_seconds
-        self.lease_seconds = lease_seconds
-        self.monitor_interval_seconds = monitor_interval_seconds
-        self._process: subprocess.Popen[bytes] | None = None
+
+    @property
+    def _cache_root(self) -> Path:
+        return self.index.project.descriptor.cache_root
+
+    @property
+    def _registry_path(self) -> Path:
+        return self._cache_root / VISUALIZATION_REGISTRY_NAME
+
+    @property
+    def _lock_path(self) -> Path:
+        return self._cache_root / VISUALIZATION_LOCK_NAME
+
+    @contextmanager
+    def _locked_registry(self) -> Generator[None, None, None]:
+        self._cache_root.mkdir(parents=True, exist_ok=True)
+        with self._lock_path.open("a+", encoding="utf-8") as handle:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
+            try:
+                yield
+            finally:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+    def _read_registry(self) -> dict[str, object] | None:
+        path = self._registry_path
+        if not path.exists():
+            return None
+        if path.is_symlink() or not path.is_file():
+            raise DocForgeError("visualization_unavailable", "Visualization registry is unsafe")
+        try:
+            raw_document: object = json.loads(path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            self._remove_registry()
+            return None
+        if not isinstance(raw_document, dict):
+            self._remove_registry()
+            return None
+        document = cast(dict[str, object], raw_document)
+        pid = document.get("pid")
+        port = document.get("port")
+        token = document.get("token")
+        snapshot = document.get("snapshot")
+        if (
+            document.get("runtime") != VISUALIZATION_RUNTIME
+            or document.get("template") != VISUALIZATION_TEMPLATE
+            or type(pid) is not int
+            or pid <= 1
+            or type(port) is not int
+            or not 1 <= port <= 65535
+            or not isinstance(token, str)
+            or len(token) < 20
+            or not isinstance(snapshot, dict)
+        ):
+            self._remove_registry()
+            return None
+        return document
+
+    def _remove_registry(self) -> None:
+        path = self._registry_path
+        if not path.exists():
+            return
+        if path.is_symlink() or not path.is_file():
+            raise DocForgeError("visualization_unavailable", "Visualization registry is unsafe")
+        path.unlink()
+
+    def _write_registry(self, document: dict[str, object]) -> None:
+        descriptor, temporary_name = tempfile.mkstemp(
+            prefix=".visualization-", dir=self._cache_root
+        )
+        temporary = Path(temporary_name)
+        try:
+            os.fchmod(descriptor, 0o600)
+            with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+                json.dump(document, handle, sort_keys=True, separators=(",", ":"))
+                handle.write("\n")
+                handle.flush()
+                os.fsync(handle.fileno())
+            os.replace(temporary, self._registry_path)
+        except OSError:
+            temporary.unlink(missing_ok=True)
+            raise
+
+    @staticmethod
+    def _worker_process(pid: int) -> bool:
+        if pid <= 1:
+            return False
+        try:
+            command = Path(f"/proc/{pid}/cmdline").read_bytes()
+        except OSError:
+            return False
+        return b"docforge.visualization_worker" in command
+
+    @staticmethod
+    def _target_url(
+        *,
+        port: int,
+        token: str,
+        node_id: str | None,
+        query: str | None,
+        depth: int,
+    ) -> str:
+        parameters: dict[str, str] = {"depth": str(depth)}
+        if node_id is not None:
+            parameters["node"] = node_id
+        if query is not None:
+            parameters["q"] = query
+        return f"http://127.0.0.1:{port}/{token}/?{urllib.parse.urlencode(parameters)}"
+
+    def _matches_snapshot(
+        self,
+        document: dict[str, object],
+        snapshot: VisualizationIndexSnapshot,
+    ) -> bool:
+        record = document.get("snapshot")
+        return isinstance(record, dict) and record == snapshot.identity
+
+    def _is_live(
+        self,
+        document: dict[str, object],
+        snapshot: VisualizationIndexSnapshot,
+    ) -> bool:
+        if not self._matches_snapshot(document, snapshot):
+            return False
+        pid = cast(int, document["pid"])
+        port = cast(int, document["port"])
+        token = cast(str, document["token"])
+        if not self._worker_process(pid):
+            return False
+        request = urllib.request.Request(
+            f"http://127.0.0.1:{port}/{token}/api/overview",
+            headers={"Accept": "application/json"},
+        )
+        try:
+            with urllib.request.urlopen(request, timeout=1) as response:
+                raw_payload: object = json.load(response)
+        except (OSError, ValueError, urllib.error.URLError):
+            return False
+        if not isinstance(raw_payload, dict):
+            return False
+        payload = cast(dict[str, object], raw_payload)
+        return all(payload.get(key) == value for key, value in snapshot.identity.items())
+
+    @staticmethod
+    def _terminate_worker(pid: int) -> None:
+        if not PersistentVisualizationRunner._worker_process(pid):
+            return
+        try:
+            os.kill(pid, signal.SIGTERM)
+        except ProcessLookupError:
+            return
+        except PermissionError as error:
+            raise DocForgeError(
+                "visualization_unavailable", "Visualization worker cannot be stopped"
+            ) from error
+        deadline = monotonic() + 2
+        while PersistentVisualizationRunner._worker_process(pid) and monotonic() < deadline:
+            time.sleep(0.05)
+        if PersistentVisualizationRunner._worker_process(pid):
+            try:
+                os.kill(pid, signal.SIGKILL)
+            except ProcessLookupError:
+                return
+
+    def _result(
+        self,
+        snapshot: VisualizationIndexSnapshot,
+        *,
+        port: int,
+        token: str,
+        node_id: str | None,
+        query: str | None,
+        depth: int,
+        reused: bool,
+    ) -> dict[str, object]:
+        return {
+            "state": "running",
+            "reused": reused,
+            "url": self._target_url(
+                port=port,
+                token=token,
+                node_id=node_id,
+                query=query,
+                depth=depth,
+            ),
+            "bind": "127.0.0.1",
+            "port": port,
+            "template": VISUALIZATION_TEMPLATE,
+            "read_only": True,
+            "project_bound": True,
+            "lifetime": {
+                "policy": "explicit_stop",
+                "stop_tool": "docforge_stop_visualization",
+            },
+            "target": {"node_id": node_id, "query": query, "depth": depth},
+            "snapshot": dict(snapshot.identity),
+        }
 
     def start(
         self,
@@ -774,8 +975,34 @@ class DetachedVisualizationRunner:
         elif query is not None:
             snapshot.search(query=query, family=None, limit=1)
 
-        self.stop()
+        with self._locked_registry():
+            existing = self._read_registry()
+            if existing is not None and self._is_live(existing, snapshot):
+                return self._result(
+                    snapshot,
+                    port=cast(int, existing["port"]),
+                    token=cast(str, existing["token"]),
+                    node_id=node_id,
+                    query=query,
+                    depth=depth,
+                    reused=True,
+                )
+            if existing is not None:
+                self._terminate_worker(cast(int, existing["pid"]))
+                self._remove_registry()
+            return self._launch(snapshot, node_id=node_id, query=query, depth=depth)
+
+    def _launch(
+        self,
+        snapshot: VisualizationIndexSnapshot,
+        *,
+        node_id: str | None,
+        query: str | None,
+        depth: int,
+    ) -> dict[str, object]:
         parent_socket, child_socket = socket.socketpair()
+        process_id: int | None = None
+        token = secrets.token_urlsafe(24)
         try:
             command = (
                 sys.executable,
@@ -784,64 +1011,80 @@ class DetachedVisualizationRunner:
                 "--control-fd",
                 str(child_socket.fileno()),
             )
-            process = subprocess.Popen(
+            child_socket.set_inheritable(True)
+            process_id = os.posix_spawn(
+                sys.executable,
                 command,
-                stdin=subprocess.DEVNULL,
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL,
-                close_fds=True,
-                pass_fds=(child_socket.fileno(),),
-                start_new_session=True,
+                os.environ,
+                setsid=True,
             )
-            self._process = process
             child_socket.close()
             request = {
                 "snapshot": snapshot.spec(),
-                "token": secrets.token_urlsafe(24),
-                "initial_grace_seconds": self.initial_grace_seconds,
-                "lease_seconds": self.lease_seconds,
-                "monitor_interval_seconds": self.monitor_interval_seconds,
-                "owner_pid": self.owner_pid,
-                "target": {
-                    "node_id": node_id,
-                    "query": query,
-                    "depth": depth,
-                },
+                "token": token,
+                "target": {"node_id": node_id, "query": query, "depth": depth},
             }
             parent_socket.settimeout(10)
             parent_socket.sendall(
                 json.dumps(request, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n"
             )
             response = _receive_worker_response(parent_socket)
-        except (OSError, subprocess.SubprocessError, ValueError) as error:
-            self.stop()
+            if response.get("status") != "ok" or not isinstance(
+                response.get("visualization"), dict
+            ):
+                raise ValueError("Visualization worker rejected the snapshot")
+            visualization = cast(dict[str, object], response["visualization"])
+            port = visualization.get("port")
+            if type(port) is not int or not 1 <= port <= 65535:
+                raise ValueError("Visualization worker returned an invalid port")
+            self._write_registry(
+                {
+                    "runtime": VISUALIZATION_RUNTIME,
+                    "template": VISUALIZATION_TEMPLATE,
+                    "pid": process_id,
+                    "port": port,
+                    "token": token,
+                    "snapshot": dict(snapshot.identity),
+                }
+            )
+            return self._result(
+                snapshot,
+                port=port,
+                token=token,
+                node_id=node_id,
+                query=query,
+                depth=depth,
+                reused=False,
+            )
+        except (OSError, ValueError) as error:
+            if process_id is not None:
+                self._terminate_worker(process_id)
             raise DocForgeError(
                 "visualization_unavailable",
-                "The detached visualization worker failed to start",
+                "The persistent visualization worker failed to start",
             ) from error
         finally:
             child_socket.close()
             parent_socket.close()
 
-        if response.get("status") != "ok" or not isinstance(response.get("visualization"), dict):
-            self.stop()
-            raise DocForgeError(
-                "visualization_unavailable",
-                "The detached visualization worker rejected the snapshot",
-            )
-        return cast(dict[str, object], response["visualization"])
+    def stop(self) -> dict[str, object]:
+        with self._locked_registry():
+            existing = self._read_registry()
+            if existing is None:
+                return self._stop_result("not_running")
+            self._terminate_worker(cast(int, existing["pid"]))
+            self._remove_registry()
+            return self._stop_result("stopped")
 
-    def stop(self) -> None:
-        process = self._process
-        self._process = None
-        if process is None or process.poll() is not None:
-            return
-        process.terminate()
-        try:
-            process.wait(timeout=2)
-        except subprocess.TimeoutExpired:
-            process.kill()
-            process.wait(timeout=2)
+    def _stop_result(self, state: str) -> dict[str, object]:
+        descriptor = self.index.project.descriptor
+        return {
+            "status": "ok",
+            "state": state,
+            "project_id": descriptor.project_id,
+            "project_root_fingerprint": project_root_fingerprint(descriptor.root),
+            "adapter": descriptor.adapter,
+        }
 
 
 def _receive_worker_response(control: socket.socket) -> dict[str, object]: