guardian_client.py

"""Small named-pipe client used by Explorer launchers and management tools."""
from __future__ import annotations

import subprocess
import time
from multiprocessing.connection import Client
from typing import Any

from guardian_common import (
    BROKER_PATH,
    PIPE_ADDRESS,
    RUNTIME_PYTHONW,
    auth_key,
)


class GuardianUnavailable(RuntimeError):
    pass


def request(message: dict[str, Any]) -> dict[str, Any]:
    connection = Client(PIPE_ADDRESS, family="AF_PIPE", authkey=auth_key())
    try:
        connection.send(message)
        response = connection.recv()
    finally:
        connection.close()
    if not isinstance(response, dict):
        raise GuardianUnavailable("Guardian returned an invalid response")
    return response


def ping() -> dict[str, Any] | None:
    try:
        return request({"action": "ping"})
    except (OSError, EOFError, ConnectionError, GuardianUnavailable):
        return None


def ensure_broker(timeout: float = 15.0) -> dict[str, Any]:
    status = ping()
    if status:
        return status
    if not RUNTIME_PYTHONW.is_file():
        raise GuardianUnavailable(f"Guardian runtime is missing: {RUNTIME_PYTHONW}")
    subprocess.Popen(
        [str(RUNTIME_PYTHONW), str(BROKER_PATH)],
        close_fds=True,
        creationflags=subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP,
    )
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        time.sleep(0.2)
        status = ping()
        if status:
            return status
    raise GuardianUnavailable("Python Guardian did not become ready")


def preflight(script: str) -> dict[str, Any]:
    ensure_broker()
    return request({"action": "preflight", "script": script})


def submit(script: str, interpreter: str, arguments: list[str], keep: str) -> dict[str, Any]:
    ensure_broker()
    return request(
        {
            "action": "submit",
            "script": script,
            "interpreter": interpreter,
            "arguments": arguments,
            "keep": keep,
        }
    )


def stop_broker() -> dict[str, Any]:
    response = request({"action": "stop"})
    # Wake Listener.accept() once more so its loop can observe stop_event.
    try:
        request({"action": "ping"})
    except (OSError, EOFError, ConnectionError, GuardianUnavailable):
        pass
    return response


def job_status(job_id: str) -> dict[str, Any]:
    return request({"action": "job", "job_id": job_id})


def start_repair(job_id: str) -> dict[str, Any]:
    ensure_broker()
    return request({"action": "repair", "job_id": job_id})


def repair_status(job_id: str) -> dict[str, Any]:
    return request({"action": "repair_status", "job_id": job_id})


def retry_job(job_id: str) -> dict[str, Any]:
    return request({"action": "retry", "job_id": job_id})


def undo_repair(job_id: str) -> dict[str, Any]:
    return request({"action": "undo_repair", "job_id": job_id})