Manage graph viewers with a supervised local service
This commit is contained in:
parent
eb48ba1a51
commit
7c87536167
14 changed files with 1027 additions and 103 deletions
755
src/docforge/viewer_manager.py
Normal file
755
src/docforge/viewer_manager.py
Normal file
|
|
@ -0,0 +1,755 @@
|
|||
"""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())
|
||||
Loading…
Add table
Add a link
Reference in a new issue