755 lines
28 KiB
Python
755 lines
28 KiB
Python
"""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, 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())
|