Detach graph viewer from MCP transport

src/docforge/visualization_worker.py

@@ -0,0 +1,110 @@
+"""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 Any
+
+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, Any]:
+    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 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"]
+        if not isinstance(target, dict):
+            raise ValueError("Visualization target is invalid")
+        runner = VisualizationRunner(
+            None,
+            snapshot_spec=request["snapshot"],
+            token=str(request["token"]),
+            register_atexit=False,
+            initial_grace_seconds=float(request["initial_grace_seconds"]),
+            lease_seconds=float(request["lease_seconds"]),
+            monitor_interval_seconds=float(request["monitor_interval_seconds"]),
+        )
+        visualization = runner.start(
+            node_id=target.get("node_id"),
+            query=target.get("query"),
+            depth=int(target["depth"]),
+        )
+        control.sendall(
+            json.dumps(
+                {"status": "ok", "visualization": visualization},
+                sort_keys=True,
+                separators=(",", ":"),
+            ).encode("utf-8")
+            + b"\n"
+        )
+        owner_pid = int(request["owner_pid"])
+    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())