aps-agent/server/sidecar.py

"""Desktop Python sidecar entry with a fail-closed native runtime preflight."""

from __future__ import annotations

import argparse
import hmac
import importlib.util
import ipaddress
import json
import os
import platform
import re
import site
import subprocess
import sys
import threading
import time
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

PROBE_OK_PREFIX = "APS_NATIVE_PROBE_OK"
SIDECAR_NONCE_HEADER = b"x-aps-sidecar-nonce"
SIDECAR_NONCE_PATTERN = re.compile(r"^[a-f0-9]{64}$")
FATAL_MARKERS = (
    "windows fatal exception",
    "fatal python error",
    "0xc0000139",
    "winerror 127",
    "access violation",
    "segmentation fault",
)


class SidecarStartupError(RuntimeError):
    """The sidecar cannot start safely with the current runtime or configuration."""


@dataclass(frozen=True)
class SidecarConfig:
    host: str
    port: int
    ui_dir: str | None
    nonce: str | None
    parent_pid: int | None


@dataclass(frozen=True)
class NativeProbeResult:
    stdout: str
    stderr: str


def normalize_loopback_host(raw_host: str) -> str:
    """Accept only explicit loopback targets; never resolve arbitrary hostnames."""
    host = raw_host.strip()
    if not host:
        raise SidecarStartupError("APS_API_HOST cannot be empty")
    if host.lower() == "localhost":
        return "127.0.0.1"
    try:
        address = ipaddress.ip_address(host)
    except ValueError as exc:
        raise SidecarStartupError(
            "APS_API_HOST must be localhost or a loopback IP address"
        ) from exc
    if not address.is_loopback:
        raise SidecarStartupError("APS_API_HOST must be loopback-only")
    return address.compressed


def parse_api_port(raw_port: str) -> int:
    try:
        port = int(raw_port, 10)
    except ValueError as exc:
        raise SidecarStartupError("APS_API_PORT must be an integer") from exc
    if not 1 <= port <= 65535:
        raise SidecarStartupError("APS_API_PORT must be between 1 and 65535")
    return port


def normalize_nonce(raw_nonce: str) -> str | None:
    nonce = raw_nonce.strip().lower()
    if not nonce:
        return None
    if SIDECAR_NONCE_PATTERN.fullmatch(nonce) is None:
        raise SidecarStartupError(
            "APS_SIDECAR_NONCE must be 64 lowercase hexadecimal characters"
        )
    return nonce


def parse_parent_pid(raw_pid: str) -> int | None:
    value = raw_pid.strip()
    if not value:
        return None
    try:
        pid = int(value, 10)
    except ValueError as exc:
        raise SidecarStartupError("APS_PARENT_PID must be a positive integer") from exc
    if pid <= 0 or pid == os.getpid():
        raise SidecarStartupError(
            "APS_PARENT_PID must identify a different positive process"
        )
    return pid


def load_config(env: Mapping[str, str] | None = None) -> SidecarConfig:
    source = os.environ if env is None else env
    return SidecarConfig(
        host=normalize_loopback_host(source.get("APS_API_HOST", "127.0.0.1")),
        port=parse_api_port(source.get("APS_API_PORT", "8000")),
        ui_dir=source.get("APS_UI_DIR", "").strip() or None,
        nonce=normalize_nonce(source.get("APS_SIDECAR_NONCE", "")),
        parent_pid=parse_parent_pid(source.get("APS_PARENT_PID", "")),
    )


def parent_is_alive(pid: int) -> bool:
    if os.name == "nt":
        import ctypes

        process_query_limited_information = 0x1000
        still_active = 259
        handle = ctypes.windll.kernel32.OpenProcess(
            process_query_limited_information, False, pid
        )
        if not handle:
            return False
        try:
            exit_code = ctypes.c_ulong()
            if not ctypes.windll.kernel32.GetExitCodeProcess(
                handle, ctypes.byref(exit_code)
            ):
                return False
            return exit_code.value == still_active
        finally:
            ctypes.windll.kernel32.CloseHandle(handle)
    try:
        os.kill(pid, 0)
    except OSError:
        return False
    return True


def watch_parent(
    parent_pid: int,
    *,
    is_alive: Callable[[int], bool] = parent_is_alive,
    exit_process: Callable[[int], Any] = os._exit,
    interval_seconds: float = 1.0,
) -> None:
    while is_alive(parent_pid):
        time.sleep(interval_seconds)
    exit_process(0)


def start_parent_watchdog(parent_pid: int) -> threading.Thread:
    watcher = threading.Thread(
        target=watch_parent,
        args=(parent_pid,),
        name="aps-parent-watchdog",
        daemon=True,
    )
    watcher.start()
    return watcher


def _module_origin(package: str) -> str | None:
    spec = importlib.util.find_spec(package)
    if spec is None:
        return None
    if spec.origin and spec.origin not in {"built-in", "frozen"}:
        return str(Path(spec.origin).resolve())
    locations = tuple(spec.submodule_search_locations or ())
    return str(Path(locations[0]).resolve()) if locations else spec.origin


def _runtime_identity(package_origins: Mapping[str, str | None]) -> dict[str, Any]:
    return {
        "runtimeIdentity": (
            f"{platform.python_implementation()} {platform.python_version()} "
            f"{platform.system()}-{platform.machine()}"
        ),
        "executable": str(Path(sys.executable).resolve()),
        "baseExecutable": str(
            Path(getattr(sys, "_base_executable", None) or sys.executable).resolve()
        ),
        "prefix": str(Path(sys.prefix).resolve()),
        "basePrefix": str(Path(sys.base_prefix).resolve()),
        "ENABLE_USER_SITE": bool(site.ENABLE_USER_SITE),
        "frozen": bool(getattr(sys, "frozen", False)),
        "packageOrigins": dict(package_origins),
    }


def _runtime_safety_errors(identity: Mapping[str, Any]) -> list[str]:
    if identity.get("frozen"):
        return []
    errors: list[str] = []
    if identity.get("ENABLE_USER_SITE"):
        errors.append("ENABLE_USER_SITE must be disabled")
    prefix = Path(str(identity["prefix"]))
    for package, origin in dict(identity.get("packageOrigins") or {}).items():
        if not origin:
            errors.append(f"{package} is not installed")
            continue
        try:
            Path(str(origin)).resolve().relative_to(prefix)
        except ValueError:
            errors.append(f"{package} is outside runtime prefix: {origin}")
    return errors


def probe_native_dependencies() -> None:
    """Import native stacks and run minimal solver/Sobol calculations in the probe child."""
    package_origins = {
        package: _module_origin(package) for package in ("ortools", "numpy", "pandas", "scipy")
    }
    identity = _runtime_identity(package_origins)
    print(
        "APS_NATIVE_PROBE_RUNTIME "
        + json.dumps(identity, ensure_ascii=True, sort_keys=True),
        flush=True,
    )
    safety_errors = _runtime_safety_errors(identity)
    if safety_errors:
        raise SidecarStartupError("unsafe solver runtime: " + "; ".join(safety_errors))

    import numpy
    import ortools
    import pandas
    import scipy
    from ortools.linear_solver import pywraplp
    from ortools.sat.python import cp_model
    from scipy.stats import qmc

    model = cp_model.CpModel()
    value = model.new_int_var(0, 10, "value")
    model.maximize(value)
    solver = cp_model.CpSolver()
    status = solver.solve(model)
    if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE) or solver.value(value) != 10:
        raise SidecarStartupError("CP-SAT smoke solve failed")
    if int(numpy.dot(numpy.array([1, 2]), numpy.array([3, 4]))) != 11:
        raise SidecarStartupError("NumPy smoke calculation failed")
    if int(pandas.DataFrame({"value": [1, 2, 3]})["value"].sum()) != 6:
        raise SidecarStartupError("Pandas smoke calculation failed")
    sobol_sample = qmc.Sobol(d=2, scramble=False).random_base2(m=2)
    expected_sobol = numpy.array([
        [0.0, 0.0], [0.5, 0.5], [0.75, 0.25], [0.25, 0.75],
    ])
    if not numpy.array_equal(sobol_sample, expected_sobol):
        raise SidecarStartupError("SciPy Sobol smoke calculation failed")
    lp_solver = pywraplp.Solver.CreateSolver("GLOP")
    if lp_solver is None:
        raise SidecarStartupError("GLOP solver is unavailable")
    lp_relax = lp_solver.NumVar(0.0, lp_solver.infinity(), "relax")
    lp_cap = lp_solver.Constraint(-lp_solver.infinity(), 0.0, "cap")
    lp_cap.SetCoefficient(lp_relax, 1.0)
    lp_residual = lp_solver.NumVar(0.0, lp_solver.infinity(), "residual")
    lp_demand = lp_solver.Constraint(10.0, lp_solver.infinity(), "demand")
    lp_demand.SetCoefficient(lp_residual, 1.0)
    lp_demand.SetCoefficient(lp_relax, 1.0)
    lp_solver.Minimize(lp_residual)
    lp_status = lp_solver.Solve()
    if (
        lp_status != pywraplp.Solver.OPTIMAL
        or abs(lp_solver.Objective().Value() - 10.0) > 1e-9
        or abs(lp_cap.dual_value() + 1.0) > 1e-9
    ):
        raise SidecarStartupError("GLOP objective/dual smoke calculation failed")

    identity["packages"] = {
        "ortools": {
            "path": str(Path(ortools.__file__).resolve()),
            "version": ortools.__version__,
        },
        "numpy": {
            "path": str(Path(numpy.__file__).resolve()),
            "version": numpy.__version__,
        },
        "pandas": {
            "path": str(Path(pandas.__file__).resolve()),
            "version": pandas.__version__,
        },
        "scipy": {
            "path": str(Path(scipy.__file__).resolve()),
            "version": scipy.__version__,
        },
    }
    print(
        f"{PROBE_OK_PREFIX} numpy={numpy.__version__} "
        f"pandas={pandas.__version__} scipy={scipy.__version__} glop=GLOP "
        f"ortools={ortools.__version__} "
        f"python={platform.python_version()} machine={platform.machine()} "
        + json.dumps(identity, ensure_ascii=True, sort_keys=True),
        flush=True,
    )


def build_probe_command(
    *,
    executable: str | None = None,
    frozen: bool | None = None,
) -> list[str]:
    binary = executable or sys.executable
    is_frozen = bool(getattr(sys, "frozen", False)) if frozen is None else frozen
    if is_frozen:
        return [binary, "--probe-child"]
    return [binary, "-X", "faulthandler", "-m", "server.sidecar", "--probe-child"]


def run_native_probe(
    *,
    timeout_seconds: float = 30.0,
    runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
) -> NativeProbeResult:
    """Run native imports out-of-process so loader crashes block server startup."""
    try:
        completed = runner(
            build_probe_command(),
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout_seconds,
            check=False,
        )
    except subprocess.TimeoutExpired as exc:
        raise SidecarStartupError(
            f"native dependency probe timed out after {timeout_seconds:g}s"
        ) from exc
    output = f"{completed.stdout}\n{completed.stderr}"
    fatal_marker = next(
        (marker for marker in FATAL_MARKERS if marker in output.lower()), None
    )
    if completed.returncode != 0:
        raise SidecarStartupError(
            f"native dependency probe exited with code {completed.returncode}: {output.strip()}"
        )
    if fatal_marker is not None:
        raise SidecarStartupError(f"native dependency probe reported {fatal_marker!r}")
    if PROBE_OK_PREFIX not in completed.stdout:
        raise SidecarStartupError(
            "native dependency probe did not emit its success marker"
        )
    return NativeProbeResult(stdout=completed.stdout, stderr=completed.stderr)


class SidecarIdentityApp:
    def __init__(self, app: Any, nonce: str) -> None:
        self.app = app
        self.nonce = nonce.encode("ascii")

    async def __call__(self, scope: dict, receive: Callable, send: Callable) -> None:
        if scope.get("type") == "http":
            headers = dict(scope.get("headers", []))
            supplied_nonce = headers.get(SIDECAR_NONCE_HEADER, b"")
            if not hmac.compare_digest(supplied_nonce, self.nonce):
                body = b"Forbidden"
                await send(
                    {
                        "type": "http.response.start",
                        "status": 403,
                        "headers": [
                            (b"content-type", b"text/plain; charset=utf-8"),
                            (b"content-length", str(len(body)).encode("ascii")),
                        ],
                    }
                )
                await send({"type": "http.response.body", "body": body})
                return

        async def send_with_identity(message: dict) -> None:
            if message.get("type") == "http.response.start":
                headers = list(message.get("headers", []))
                headers.append((SIDECAR_NONCE_HEADER, self.nonce))
                message = {**message, "headers": headers}
            await send(message)

        await self.app(scope, receive, send_with_identity)


def build_desktop_app(ui_dir: str | None, nonce: str | None) -> Any:
    if ui_dir is None:
        app: Any = "server.main:app"
        if nonce is not None:
            from server.main import app as server_app

            app = server_app
        return SidecarIdentityApp(app, nonce) if nonce is not None else app
    root = Path(ui_dir).resolve()
    if not (root / "index.html").is_file():
        raise SidecarStartupError(f"desktop UI index is missing: {root}")

    from fastapi.staticfiles import StaticFiles

    from server.main import app

    app.mount("/", StaticFiles(directory=str(root), html=True), name="desktop-ui")
    return SidecarIdentityApp(app, nonce) if nonce is not None else app


def run_sidecar(
    *,
    env: Mapping[str, str] | None = None,
    probe: Callable[[], NativeProbeResult] = run_native_probe,
    uvicorn_runner: Callable[..., Any] | None = None,
    app_builder: Callable[[str | None, str | None], Any] = build_desktop_app,
    parent_watchdog: Callable[[int], Any] = start_parent_watchdog,
) -> int:
    """Validate configuration, probe native dependencies, then block in uvicorn."""
    config = load_config(env)
    probe()
    os.environ.setdefault("APS_MODE", "desktop")
    if config.parent_pid is not None:
        parent_watchdog(config.parent_pid)
    if uvicorn_runner is None:
        import uvicorn

        uvicorn_runner = uvicorn.run
    uvicorn_runner(
        app_builder(config.ui_dir, config.nonce),
        host=config.host,
        port=config.port,
        log_level="info",
        workers=1,
    )
    return 0


def run_solver_child(argv: Sequence[str] | None = None) -> int:
    """Run the W66 stdin/stdout worker protocol without starting the Web app."""
    if not getattr(sys, "frozen", False):
        source_root = Path(__file__).resolve().parents[1]
        worker_source = source_root / "server" / "engines" / "solver_worker.py"
        source_root_text = str(source_root)
        if worker_source.is_file() and source_root_text not in sys.path:
            # ``python -I server/sidecar.py`` intentionally removes cwd/script
            # import roots. Re-add only this resolved, trusted source checkout so
            # the isolated child can import the packaged ``server`` namespace.
            sys.path.insert(0, source_root_text)
    try:
        from server.engines.solver_worker import main as solver_worker_main
    except (ImportError, AttributeError) as exc:
        raise SidecarStartupError("solver worker protocol is unavailable") from exc
    previous_argv = sys.argv[:]
    sys.argv = ["server.engines.solver_worker", *list(argv or ())]
    try:
        result = solver_worker_main()
    finally:
        sys.argv = previous_argv
    return 0 if result is None else int(result)


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run the APS desktop Python sidecar.")
    child_mode = parser.add_mutually_exclusive_group()
    child_mode.add_argument(
        "--probe-child", action="store_true", help=argparse.SUPPRESS
    )
    child_mode.add_argument(
        "--solver-child", action="store_true", help=argparse.SUPPRESS
    )
    args, child_args = parser.parse_known_args(argv)
    if child_args and not args.solver_child:
        parser.error(f"unrecognized arguments: {' '.join(child_args)}")
    try:
        if args.probe_child:
            probe_native_dependencies()
            return 0
        if args.solver_child:
            return run_solver_child(child_args)
        return run_sidecar()
    except SidecarStartupError as exc:
        print(f"APS_SIDECAR_STARTUP_ERROR: {exc}", file=sys.stderr, flush=True)
        return 2


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