"""Detached process host for one project-bound DocForge visualization.""" from __future__ import annotations import argparse import json import os import socket import time from contextlib import suppress from typing import cast from .errors import DocForgeError from .visualization import VisualizationRunner def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="docforge-visualization-worker") parser.add_argument("--control-fd", type=int, required=True) return parser def _read_request(control: socket.socket) -> dict[str, object]: chunks: list[bytes] = [] size = 0 while True: chunk = control.recv(65536) if not chunk: break chunks.append(chunk) size += len(chunk) if size > 1_000_000: raise ValueError("Visualization launch request exceeded its fixed boundary") if b"\n" in chunk: break payload = b"".join(chunks).split(b"\n", 1)[0] request = json.loads(payload) if not isinstance(request, dict): raise ValueError("Visualization launch request is invalid") return cast(dict[str, object], request) def _pid_exists(pid: int) -> bool: if pid <= 1: return False try: os.kill(pid, 0) except ProcessLookupError: return False except PermissionError: return True return True def main(argv: list[str] | None = None) -> int: arguments = _parser().parse_args(argv) control = socket.socket(fileno=arguments.control_fd) runner: VisualizationRunner | None = None try: request = _read_request(control) target = request["target"] snapshot = request["snapshot"] token = request["token"] initial_grace = request["initial_grace_seconds"] lease = request["lease_seconds"] monitor_interval = request["monitor_interval_seconds"] owner = request["owner_pid"] if not isinstance(target, dict) or not isinstance(snapshot, dict): raise ValueError("Visualization target is invalid") 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 or not isinstance(token, str) or not isinstance(initial_grace, int | float) or isinstance(initial_grace, bool) or not isinstance(lease, int | float) or isinstance(lease, bool) or not isinstance(monitor_interval, int | float) or isinstance(monitor_interval, bool) or type(owner) is not int ): raise ValueError("Visualization launch request is invalid") runner = VisualizationRunner( None, snapshot_spec=cast(dict[str, object], snapshot), token=token, register_atexit=False, initial_grace_seconds=float(initial_grace), lease_seconds=float(lease), monitor_interval_seconds=float(monitor_interval), ) visualization = runner.start( node_id=node_id, query=query, depth=depth, ) control.sendall( json.dumps( {"status": "ok", "visualization": visualization}, sort_keys=True, separators=(",", ":"), ).encode("utf-8") + b"\n" ) owner_pid = owner except (DocForgeError, KeyError, OSError, TypeError, ValueError) as error: with suppress(OSError): control.sendall( json.dumps( {"status": "error", "error": type(error).__name__}, sort_keys=True, separators=(",", ":"), ).encode("utf-8") + b"\n" ) if runner is not None: runner.stop() return 2 finally: control.close() while runner.is_running() and _pid_exists(owner_pid): time.sleep(0.25) runner.stop() return 0 if __name__ == "__main__": raise SystemExit(main())