DocForge2/src/docforge/visualization_worker.py

"""Detached process host for one project-bound DocForge visualization."""

from __future__ import annotations

import argparse
import json
import socket
import sys
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")
    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 | 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:
        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 _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 = None if arguments.request_stdin else socket.socket(fileno=arguments.control_fd)
    runner: VisualizationRunner | None = None
    try:
        request = _read_request(control)
        target = request["target"]
        snapshot = request["snapshot"]
        token = request["token"]
        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)
        ):
            raise ValueError("Visualization launch request is invalid")
        runner = VisualizationRunner(
            None,
            snapshot_spec=cast(dict[str, object], snapshot),
            token=token,
            register_atexit=False,
            persistent=True,
        )
        visualization = runner.start(
            node_id=node_id,
            query=query,
            depth=depth,
        )
        _send_response(control, {"status": "ok", "visualization": visualization})
    except (DocForgeError, KeyError, OSError, TypeError, ValueError) as error:
        with suppress(OSError):
            _send_response(control, {"status": "error", "error": type(error).__name__})
        if runner is not None:
            runner.stop()
        return 2
    finally:
        if control is not None:
            control.close()

    while runner.is_running():
        time.sleep(0.25)
    runner.stop()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())