"""Cross-platform local manager for persistent DocForge graph viewers.""" from __future__ import annotations import argparse import json import os import plistlib import secrets import shlex import signal import socket import subprocess import sys import tempfile import threading import time import urllib.error import urllib.parse import urllib.request from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from typing import BinaryIO, cast from .errors import DocForgeError from .index import ProjectIndex from .project import project_root_fingerprint from .visualization import VISUALIZATION_TEMPLATE, VisualizationIndexSnapshot MANAGER_PROTOCOL = "docforge-viewer-manager@1" MANAGER_RUNTIME = "viewer-manager@1" DEFAULT_IDLE_TIMEOUT_SECONDS = 3600.0 DEFAULT_CHECK_INTERVAL_SECONDS = 30.0 MAX_MESSAGE_BYTES = 1_000_000 def default_runtime_root() -> Path: """Return the current user's private state directory on each supported OS.""" configured = os.environ.get("DOCFORGE_VIEWER_MANAGER_HOME") if configured: return Path(configured) if sys.platform == "win32": return Path(os.environ.get("LOCALAPPDATA", Path.home())) / "DocForge" if sys.platform == "darwin": return Path.home() / "Library" / "Application Support" / "DocForge" runtime = os.environ.get("XDG_RUNTIME_DIR") if runtime: return Path(runtime) / "docforge" return Path.home() / ".cache" / "docforge" def default_state_path() -> Path: configured = os.environ.get("DOCFORGE_VIEWER_MANAGER_STATE") return Path(configured) if configured else default_runtime_root() / "viewer-manager.json" def user_service_path() -> Path | None: """Return the OS-specific user-service definition path, when one is file-backed.""" if sys.platform.startswith("linux"): return Path.home() / ".config" / "systemd" / "user" / "docforge-viewer-manager.service" if sys.platform == "darwin": return Path.home() / "Library" / "LaunchAgents" / "com.docforge.viewer-manager.plist" return None def _manager_command(state_path: Path) -> list[str]: return [ sys.executable, "-m", "docforge.viewer_manager", "serve", "--state-path", str(state_path), ] def _systemd_unit(state_path: Path) -> str: command = " ".join(shlex.quote(argument) for argument in _manager_command(state_path)) return ( "[Unit]\n" "Description=DocForge graph viewer manager\n" "\n" "[Service]\n" "Type=simple\n" f"ExecStart={command}\n" "Restart=on-failure\n" "RestartSec=2\n" "KillMode=control-group\n" "NoNewPrivileges=true\n" "\n" "[Install]\n" "WantedBy=default.target\n" ) def _launch_agent(state_path: Path) -> bytes: return plistlib.dumps( { "Label": "com.docforge.viewer-manager", "ProgramArguments": _manager_command(state_path), "RunAtLoad": True, "KeepAlive": {"SuccessfulExit": False}, "ProcessType": "Background", }, fmt=plistlib.FMT_XML, sort_keys=False, ) def install_user_service(state_path: Path) -> str: """Install and start the native per-user supervisor for this platform.""" if sys.platform.startswith("linux"): path = user_service_path() assert path is not None path.parent.mkdir(parents=True, exist_ok=True) path.write_text(_systemd_unit(state_path), encoding="utf-8") subprocess.run(["systemctl", "--user", "daemon-reload"], check=True) subprocess.run(["systemctl", "--user", "enable", "--now", path.name], check=True) return str(path) if sys.platform == "darwin": path = user_service_path() assert path is not None path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(_launch_agent(state_path)) domain = f"gui/{os.getuid()}" subprocess.run(["launchctl", "bootout", domain, str(path)], check=False) subprocess.run(["launchctl", "bootstrap", domain, str(path)], check=True) subprocess.run( ["launchctl", "kickstart", "-k", f"{domain}/com.docforge.viewer-manager"], check=True, ) return str(path) if sys.platform == "win32": command = subprocess.list2cmdline(_manager_command(state_path)) subprocess.run( [ "schtasks", "/create", "/tn", "DocForgeViewerManager", "/tr", command, "/sc", "onlogon", "/rl", "limited", "/f", ], check=True, ) subprocess.run(["schtasks", "/run", "/tn", "DocForgeViewerManager"], check=True) return "Task Scheduler: DocForgeViewerManager" raise RuntimeError(f"No user-service installer is available for {sys.platform}") def uninstall_user_service() -> None: """Stop and remove the native per-user viewer-manager supervisor.""" if sys.platform.startswith("linux"): path = user_service_path() assert path is not None subprocess.run(["systemctl", "--user", "disable", "--now", path.name], check=False) path.unlink(missing_ok=True) subprocess.run(["systemctl", "--user", "daemon-reload"], check=False) return if sys.platform == "darwin": path = user_service_path() assert path is not None subprocess.run(["launchctl", "bootout", f"gui/{os.getuid()}", str(path)], check=False) path.unlink(missing_ok=True) return if sys.platform == "win32": subprocess.run(["schtasks", "/delete", "/tn", "DocForgeViewerManager", "/f"], check=False) return raise RuntimeError(f"No user-service installer is available for {sys.platform}") def _read_message(stream: BinaryIO) -> dict[str, object]: line = stream.readline(MAX_MESSAGE_BYTES + 1) if not line or len(line) > MAX_MESSAGE_BYTES: raise ValueError("Viewer-manager message is invalid") decoded: object = json.loads(line) if not isinstance(decoded, dict): raise ValueError("Viewer-manager message is invalid") return cast(dict[str, object], decoded) def _write_message(stream: BinaryIO, payload: Mapping[str, object]) -> None: stream.write(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n") stream.flush() def _project_key(identity: dict[str, object]) -> str: project_id = identity.get("project_id") fingerprint = identity.get("project_root_fingerprint") if not isinstance(project_id, str) or not project_id: raise ValueError("Viewer-manager project identity is invalid") if not isinstance(fingerprint, str) or not fingerprint: raise ValueError("Viewer-manager project identity is invalid") return f"{project_id}:{fingerprint}" 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)}" @dataclass class _ManagedWorker: process: subprocess.Popen[bytes] port: int token: str snapshot: dict[str, object] last_activity_at: float class ViewerManager: """Own project viewer processes behind an authenticated loopback control API.""" def __init__( self, state_path: Path, *, idle_timeout_seconds: float = DEFAULT_IDLE_TIMEOUT_SECONDS, check_interval_seconds: float = DEFAULT_CHECK_INTERVAL_SECONDS, ) -> None: if idle_timeout_seconds <= 0 or check_interval_seconds <= 0: raise ValueError("Viewer-manager durations must be positive") self.state_path = state_path self.idle_timeout_seconds = idle_timeout_seconds self.check_interval_seconds = check_interval_seconds self._token = secrets.token_urlsafe(32) self._workers: dict[str, _ManagedWorker] = {} self._lock = threading.RLock() self._stopping = threading.Event() self._server: socket.socket | None = None def serve_forever(self) -> None: self._prepare_state_path() server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0) server.bind(("127.0.0.1", 0)) server.listen(32) server.settimeout(0.5) self._server = server port = cast(int, server.getsockname()[1]) self._write_state(port) monitor = threading.Thread( target=self._monitor_workers, name="docforge-viewer-manager-monitor", daemon=True, ) monitor.start() try: while not self._stopping.is_set(): try: connection, _ = server.accept() except TimeoutError: continue except OSError: if self._stopping.is_set(): break raise thread = threading.Thread( target=self._serve_connection, args=(connection,), name="docforge-viewer-manager-request", daemon=True, ) thread.start() finally: self.shutdown() monitor.join(timeout=2) def shutdown(self) -> None: if self._stopping.is_set(): return self._stopping.set() server = self._server self._server = None if server is not None: server.close() with self._lock: workers = list(self._workers.values()) self._workers.clear() for worker in workers: self._stop_worker(worker) self._remove_state() def _prepare_state_path(self) -> None: directory = self.state_path.parent directory.mkdir(parents=True, exist_ok=True) if os.name != "nt": os.chmod(directory, 0o700) if not self.state_path.exists(): return if self.state_path.is_symlink() or not self.state_path.is_file(): raise RuntimeError("Viewer-manager state path is unsafe") self.state_path.unlink() def _write_state(self, port: int) -> None: descriptor, temporary_name = tempfile.mkstemp( prefix=".viewer-manager-", dir=self.state_path.parent ) temporary = Path(temporary_name) try: if os.name != "nt": os.fchmod(descriptor, 0o600) with os.fdopen(descriptor, "w", encoding="utf-8") as stream: json.dump( { "runtime": MANAGER_RUNTIME, "host": "127.0.0.1", "port": port, "token": self._token, }, stream, sort_keys=True, separators=(",", ":"), ) stream.write("\n") stream.flush() os.fsync(stream.fileno()) os.replace(temporary, self.state_path) except OSError: temporary.unlink(missing_ok=True) raise def _remove_state(self) -> None: if not self.state_path.exists(): return if self.state_path.is_symlink() or not self.state_path.is_file(): return self.state_path.unlink() def _serve_connection(self, connection: socket.socket) -> None: with connection: connection.settimeout(10) reader = connection.makefile("rb") writer = connection.makefile("wb") try: request = _read_message(reader) response = self._handle(request) except (DocForgeError, OSError, TypeError, ValueError) as error: response = {"status": "error", "error": type(error).__name__} try: _write_message(writer, response) finally: writer.close() reader.close() def _handle(self, request: dict[str, object]) -> dict[str, object]: if request.get("protocol") != MANAGER_PROTOCOL or request.get("token") != self._token: raise ValueError("Viewer-manager authentication failed") action = request.get("action") if action == "start": return self._start(request) if action == "stop": return self._stop(request) if action == "status": return self._status(request) raise ValueError("Viewer-manager action is unsupported") def _start(self, request: dict[str, object]) -> dict[str, object]: snapshot_spec = request.get("snapshot") target = request.get("target") if not isinstance(snapshot_spec, dict) or not isinstance(target, dict): raise ValueError("Viewer-manager start request is invalid") snapshot = VisualizationIndexSnapshot.from_spec(cast(dict[str, object], snapshot_spec)) target = cast(dict[str, object], target) node_id = target.get("node_id") query = target.get("query") depth = target.get("depth") if ( (node_id is not None and not isinstance(node_id, str)) or (query is not None and not isinstance(query, str)) or type(depth) is not int ): raise ValueError("Viewer-manager target is invalid") key = _project_key(snapshot.identity) with self._lock: existing = self._workers.get(key) if existing is not None and self._is_current(existing, snapshot): return { "status": "ok", "visualization": self._result(existing, node_id, query, depth, True), } if existing is not None: self._workers.pop(key, None) self._stop_worker(existing) worker = self._launch(snapshot, node_id, query, depth) self._workers[key] = worker return { "status": "ok", "visualization": self._result(worker, node_id, query, depth, False), } def _stop(self, request: dict[str, object]) -> dict[str, object]: key = self._request_key(request) with self._lock: worker = self._workers.pop(key, None) if worker is None: return {"status": "ok", "state": "not_running"} self._stop_worker(worker) return {"status": "ok", "state": "stopped"} def _status(self, request: dict[str, object]) -> dict[str, object]: key = self._request_key(request) with self._lock: worker = self._workers.get(key) if worker is None: return {"status": "ok", "state": "not_running"} activity = self._health(worker) if activity is None: self._workers.pop(key, None) self._stop_worker(worker) return {"status": "ok", "state": "not_running"} worker.last_activity_at = activity return { "status": "ok", "state": "running", "snapshot": dict(worker.snapshot), "idle_seconds": max(0, int(time.time() - activity)), "idle_timeout_seconds": self.idle_timeout_seconds, } @staticmethod def _request_key(request: dict[str, object]) -> str: identity = request.get("identity") if not isinstance(identity, dict): raise ValueError("Viewer-manager project identity is invalid") return _project_key(cast(dict[str, object], identity)) def _launch( self, snapshot: VisualizationIndexSnapshot, node_id: str | None, query: str | None, depth: int, ) -> _ManagedWorker: token = secrets.token_urlsafe(24) process = subprocess.Popen( (sys.executable, "-m", "docforge.visualization_worker", "--request-stdin"), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, ) try: if process.stdin is None or process.stdout is None: raise RuntimeError("Visualization worker control streams are unavailable") _write_message( cast(BinaryIO, process.stdin), { "snapshot": snapshot.spec(), "token": token, "target": {"node_id": node_id, "query": query, "depth": depth}, }, ) process.stdin.close() response = self._worker_response(cast(BinaryIO, process.stdout)) process.stdout.close() visualization = response.get("visualization") if response.get("status") != "ok" or not isinstance(visualization, dict): raise ValueError("Visualization worker rejected the snapshot") visualization = cast(dict[str, object], visualization) port = visualization.get("port") if type(port) is not int or not 1 <= port <= 65535: raise ValueError("Visualization worker returned an invalid port") return _ManagedWorker(process, port, token, dict(snapshot.identity), time.time()) except (OSError, RuntimeError, ValueError): self._stop_worker(_ManagedWorker(process, 0, token, {}, 0)) raise @staticmethod def _worker_response(stream: BinaryIO) -> dict[str, object]: response: dict[str, object] | None = None def read() -> None: nonlocal response response = _read_message(stream) reader = threading.Thread(target=read, name="docforge-viewer-worker-launch", daemon=True) reader.start() reader.join(timeout=10) if reader.is_alive() or response is None: raise ValueError("Visualization worker did not start in time") return response def _is_current(self, worker: _ManagedWorker, snapshot: VisualizationIndexSnapshot) -> bool: if worker.snapshot != snapshot.identity or worker.process.poll() is not None: return False activity = self._health(worker) if activity is None: return False worker.last_activity_at = activity return True @staticmethod def _stop_worker(worker: _ManagedWorker) -> None: if worker.process.poll() is not None: return worker.process.terminate() try: worker.process.wait(timeout=2) except subprocess.TimeoutExpired: worker.process.kill() worker.process.wait(timeout=2) @staticmethod def _health(worker: _ManagedWorker) -> float | None: request = urllib.request.Request( f"http://127.0.0.1:{worker.port}/{worker.token}/api/health", headers={"Accept": "application/json"}, ) try: with urllib.request.urlopen(request, timeout=1) as response: payload: object = json.load(response) except (OSError, ValueError, urllib.error.URLError): return None if not isinstance(payload, dict): return None payload = cast(dict[str, object], payload) if payload.get("viewer") != "alive": return None activity = payload.get("last_activity_at") return float(activity) if isinstance(activity, int | float) else None def _result( self, worker: _ManagedWorker, node_id: str | None, query: str | None, depth: int, reused: bool, ) -> dict[str, object]: return { "state": "running", "reused": reused, "url": _target_url( port=worker.port, token=worker.token, node_id=node_id, query=query, depth=depth, ), "bind": "127.0.0.1", "port": worker.port, "template": VISUALIZATION_TEMPLATE, "read_only": True, "project_bound": True, "lifetime": { "policy": "managed_idle", "idle_timeout_seconds": self.idle_timeout_seconds, "stop_tool": "docforge_stop_visualization", "status_tool": "docforge_visualization_status", }, "target": {"node_id": node_id, "query": query, "depth": depth}, "snapshot": dict(worker.snapshot), } def _monitor_workers(self) -> None: while not self._stopping.wait(self.check_interval_seconds): cutoff = time.time() - self.idle_timeout_seconds with self._lock: expired: list[tuple[str, _ManagedWorker]] = [] for key, worker in self._workers.items(): activity = self._health(worker) if activity is None or activity < cutoff: expired.append((key, worker)) else: worker.last_activity_at = activity for key, worker in expired: self._workers.pop(key, None) self._stop_worker(worker) class ViewerManagerClient: """Project-bound MCP-side client for the separately supervised manager service.""" def __init__(self, index: ProjectIndex, *, state_path: Path | None = None) -> None: self.index = index self.state_path = state_path or default_state_path() def start( self, *, node_id: str | None = None, query: str | None = None, depth: int = 1, ) -> dict[str, object]: if node_id is not None and query is not None: raise DocForgeError( "invalid_visualization_target", "Choose either one exact node ID or one search query", ) snapshot = VisualizationIndexSnapshot(self.index, self.index.check()) if type(depth) is not int or depth < 1 or depth > snapshot.max_depth: raise DocForgeError( "invalid_depth", "Visualization depth is outside the configured traversal limit" ) if node_id is not None: snapshot.require_node(node_id) elif query is not None: snapshot.search( query=query, family=None, kind=None, language=None, capability=None, limit=1, ) response = self._request( { "action": "start", "snapshot": snapshot.spec(), "target": {"node_id": node_id, "query": query, "depth": depth}, } ) visualization = response.get("visualization") if response.get("status") != "ok" or not isinstance(visualization, dict): raise DocForgeError("visualization_unavailable", "Viewer manager rejected the request") return cast(dict[str, object], visualization) def stop(self) -> dict[str, object]: return self._lifecycle_request("stop") def status(self) -> dict[str, object]: return self._lifecycle_request("status") def _lifecycle_request(self, action: str) -> dict[str, object]: descriptor = self.index.project.descriptor identity = { "project_id": descriptor.project_id, "project_root_fingerprint": project_root_fingerprint(descriptor.root), } response = self._request({"action": action, "identity": identity}) if response.get("status") != "ok": raise DocForgeError("visualization_unavailable", "Viewer manager rejected the request") return { **response, "project_id": descriptor.project_id, "project_root_fingerprint": identity["project_root_fingerprint"], "adapter": descriptor.adapter, } def _request(self, request: dict[str, object]) -> dict[str, object]: state = self._read_state() host = state["host"] port = state["port"] token = state["token"] connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM) try: connection.settimeout(10) connection.connect((host, port)) writer = connection.makefile("wb") reader = connection.makefile("rb") try: _write_message( writer, {"protocol": MANAGER_PROTOCOL, "token": token, **request}, ) return _read_message(reader) finally: writer.close() reader.close() except OSError as error: raise DocForgeError( "visualization_manager_unavailable", "DocForge viewer manager is not running; start its user service", state_path=str(self.state_path), ) from error finally: connection.close() def _read_state(self) -> dict[str, object]: if ( not self.state_path.exists() or self.state_path.is_symlink() or not self.state_path.is_file() ): raise DocForgeError( "visualization_manager_unavailable", "DocForge viewer manager is not running; start its user service", state_path=str(self.state_path), ) try: raw: object = json.loads(self.state_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as error: raise DocForgeError( "visualization_manager_unavailable", "Viewer-manager state is invalid" ) from error if not isinstance(raw, dict): raise DocForgeError( "visualization_manager_unavailable", "Viewer-manager state is invalid" ) state = cast(dict[str, object], raw) host = state.get("host") port = state.get("port") token = state.get("token") if ( state.get("runtime") != MANAGER_RUNTIME or host != "127.0.0.1" or type(port) is not int or not 1 <= port <= 65535 or not isinstance(token, str) or len(token) < 32 ): raise DocForgeError( "visualization_manager_unavailable", "Viewer-manager state is invalid" ) return {"host": host, "port": port, "token": token} def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="docforge-viewer-manager") parser.add_argument( "operation", choices=("serve", "install-user-service", "uninstall-user-service"), ) parser.add_argument("--state-path", type=Path, default=default_state_path()) parser.add_argument("--idle-timeout-seconds", type=float, default=DEFAULT_IDLE_TIMEOUT_SECONDS) parser.add_argument( "--check-interval-seconds", type=float, default=DEFAULT_CHECK_INTERVAL_SECONDS ) return parser def main(argv: list[str] | None = None) -> int: arguments = _parser().parse_args(argv) if arguments.operation == "install-user-service": print(install_user_service(arguments.state_path)) return 0 if arguments.operation == "uninstall-user-service": uninstall_user_service() return 0 manager = ViewerManager( arguments.state_path, idle_timeout_seconds=arguments.idle_timeout_seconds, check_interval_seconds=arguments.check_interval_seconds, ) def stop(_signal: int, _frame: object) -> None: manager.shutdown() signal.signal(signal.SIGTERM, stop) signal.signal(signal.SIGINT, stop) manager.serve_forever() return 0 if __name__ == "__main__": raise SystemExit(main())