Manage graph viewers with a supervised local service

src/docforge/visualization_worker.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import argparse
 import json
 import socket
+import sys
 import time
 from contextlib import suppress
 from typing import cast
@@ -15,11 +16,21 @@ from .visualization import VisualizationRunner
 
 def _parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(prog="docforge-visualization-worker")
-    parser.add_argument("--control-fd", type=int, required=True)
+    control = parser.add_mutually_exclusive_group(required=True)
+    control.add_argument("--control-fd", type=int)
+    control.add_argument("--request-stdin", action="store_true")
     return parser
 
 
-def _read_request(control: socket.socket) -> dict[str, object]:
+def _read_request(control: socket.socket | None) -> dict[str, object]:
+    if control is None:
+        payload = sys.stdin.buffer.readline(1_000_001)
+        if not payload or len(payload) > 1_000_000:
+            raise ValueError("Visualization launch request is invalid")
+        request: object = json.loads(payload)
+        if not isinstance(request, dict):
+            raise ValueError("Visualization launch request is invalid")
+        return cast(dict[str, object], request)
     chunks: list[bytes] = []
     size = 0
     while True:
@@ -39,9 +50,18 @@ def _read_request(control: socket.socket) -> dict[str, object]:
     return cast(dict[str, object], request)
 
 
+def _send_response(control: socket.socket | None, payload: dict[str, object]) -> None:
+    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n"
+    if control is None:
+        sys.stdout.buffer.write(encoded)
+        sys.stdout.buffer.flush()
+    else:
+        control.sendall(encoded)
+
+
 def main(argv: list[str] | None = None) -> int:
     arguments = _parser().parse_args(argv)
-    control = socket.socket(fileno=arguments.control_fd)
+    control = None if arguments.request_stdin else socket.socket(fileno=arguments.control_fd)
     runner: VisualizationRunner | None = None
     try:
         request = _read_request(control)
@@ -73,29 +93,16 @@ def main(argv: list[str] | None = None) -> int:
             query=query,
             depth=depth,
         )
-        control.sendall(
-            json.dumps(
-                {"status": "ok", "visualization": visualization},
-                sort_keys=True,
-                separators=(",", ":"),
-            ).encode("utf-8")
-            + b"\n"
-        )
+        _send_response(control, {"status": "ok", "visualization": visualization})
     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"
-            )
+            _send_response(control, {"status": "error", "error": type(error).__name__})
         if runner is not None:
             runner.stop()
         return 2
     finally:
-        control.close()
+        if control is not None:
+            control.close()
 
     while runner.is_running():
         time.sleep(0.25)