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 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 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"]
        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,
        )
        control.sendall(
            json.dumps(
                {"status": "ok", "visualization": visualization},
                sort_keys=True,
                separators=(",", ":"),
            ).encode("utf-8")
            + b"\n"
        )
    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():
        time.sleep(0.25)
    runner.stop()
    return 0


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