rune/shared/native_units.py

"""Thin adapter for native (compiled) units — R1 §5 interop contract, Wave 0.

The contract under proof: compiled units are standalone binaries exposing a
CLI boundary — argv in, JSON on stdout, non-zero exit on failure, a mandatory
``--version`` flag. This module is the Python side: run a unit, enforce a
timeout, capture stderr, and return a TYPED result. It never raises for unit
failures (those are data in the result); it raises only for adapter misuse
(missing binary, bad timeout).

No IPC daemon, no PyO3 — the file-based store stays the shared truth. If
research later surfaces a materially better pattern it competes through the
research pipeline; the CLI boundary is the default, not a cage.
"""

from __future__ import annotations

import json
import subprocess
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
PROOF_UNIT = REPO_ROOT / "interop" / "rebase-proof" / "target" / "release" / "rebase-proof"


@dataclass
class NativeResult:
    """Typed result of one native-unit invocation.

    ok           — the invocation succeeded (exit 0 AND not timed out); when
                    the unit emitted JSON with an ``ok`` field, that field must
                    not be False (plain-text output like ``--version`` has none)
    exit_code    — raw process exit code
    data         — parsed JSON object from stdout, or None if unparseable
    stderr       — captured stderr text (may be non-empty even on success)
    duration_ms  — wall-clock cost of the invocation
    timed_out    — True when the unit exceeded the timeout and was killed
    events       — in streaming mode: every parsed JSON object line, in
                    arrival order (empty for non-streaming :func:`run_native`)
    """

    ok: bool
    exit_code: int
    data: dict | None
    stderr: str
    duration_ms: int
    timed_out: bool = False
    raw_stdout: str = field(default="", repr=False)
    events: list[dict] = field(default_factory=list)


def run_native(binary: Path, args: list[str], timeout_s: float = 10.0) -> NativeResult:
    """Invoke a native unit and return a typed result.

    Raises FileNotFoundError only when *binary* itself is absent (adapter
    misuse) — every unit-level failure is reported via the result, never an
    exception, so callers can branch on ``result.ok`` without try/except.
    """
    if not binary.is_file():
        raise FileNotFoundError(f"native unit not built: {binary}")

    started = time.monotonic()
    timed_out = False
    try:
        proc = subprocess.run(
            [str(binary), *args],
            capture_output=True,
            text=True,
            timeout=timeout_s,
            check=False,  # exit codes are data, never exceptions
        )
        exit_code, stdout, stderr = proc.returncode, proc.stdout, proc.stderr
    except subprocess.TimeoutExpired as exc:
        timed_out = True
        exit_code = -1
        stdout = exc.stdout or "" if isinstance(exc.stdout, str) else (exc.stdout or b"").decode(errors="replace")
        stderr = exc.stderr or "" if isinstance(exc.stderr, str) else (exc.stderr or b"").decode(errors="replace")

    duration_ms = int((time.monotonic() - started) * 1000)

    data: dict | None = None
    stripped = stdout.strip()
    if stripped:
        try:
            parsed = json.loads(stripped.splitlines()[-1])
            if isinstance(parsed, dict):
                data = parsed
        except json.JSONDecodeError:
            data = None

    ok = (not timed_out) and exit_code == 0 and (data is None or data.get("ok") is not False)
    return NativeResult(
        ok=ok,
        exit_code=exit_code,
        data=data,
        stderr=stderr,
        duration_ms=duration_ms,
        timed_out=timed_out,
        raw_stdout=stdout,
    )


def run_proof(args: list[str], timeout_s: float = 10.0) -> NativeResult:
    """Run the Wave 0 proof unit with the standard contract defaults."""
    return run_native(PROOF_UNIT, args, timeout_s=timeout_s)


def stream_native(
    binary: Path,
    args: list[str],
    timeout_s: float = 10.0,
    on_event=None,
):
    """Run a native unit in JSONL streaming mode, yielding events as they arrive.

    The JSONL streaming contract (rubric §W7): one JSON object per stdout
    line — ``{"ok": true, "event": ...}`` progress objects along the way, and
    a terminal envelope as the FINAL line. This generator yields each parsed
    event dict the moment it arrives (so a long op surfaces partial progress
    instead of blocking until exit), then yields exactly one terminal
    :class:`NativeResult`. *on_event*, when given, is called with each parsed
    event as it arrives (the TUI wraps this in ``call_from_thread``).

    Stream-mode ok-semantics are stricter than :func:`run_native`: ``ok``
    requires exit 0, no timeout, AND a valid terminal envelope (the final
    stdout line parsed as a JSON object whose ``ok`` is not False). Malformed
    interior lines are skipped and counted in ``malformed_lines``; malformed
    FINAL line means no valid terminal envelope, so ``ok`` is False — partial
    progress stays available in ``events``. That is the point of streaming.

    Same failure discipline as :func:`run_native`: raises FileNotFoundError
    when *binary* is absent; every unit-level failure (bad envelope, timeout,
    non-zero exit) is data on the terminal result, never an exception.
    """
    if on_event is not None and not callable(on_event):
        raise TypeError("on_event must be callable")
    if not binary.is_file():
        raise FileNotFoundError(f"native unit not built: {binary}")

    started = time.monotonic()
    proc = subprocess.Popen(
        [str(binary), *args],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )

    # stderr must drain concurrently — reading it after the stdout loop
    # deadlocks when a unit floods stderr past the pipe buffer mid-stream.
    stderr_parts: list[str] = []
    stderr_drain = threading.Thread(
        target=lambda: stderr_parts.append(proc.stderr.read()), daemon=True
    )
    stderr_drain.start()

    killed = threading.Event()
    # The watchdog's kill and the flag are one step so an adapter-side
    # timeout is never mistaken for a unit dying by its own signal.
    watchdog = threading.Timer(timeout_s, lambda: (killed.set(), proc.kill()))
    watchdog.start()

    events: list[dict] = []
    malformed_lines = 0
    # True only while the most recent non-empty stdout line parsed as a JSON
    # object — a malformed final line invalidates the terminal envelope.
    final_line_is_json = False
    try:
        for line in proc.stdout:
            line = line.strip()
            if not line:
                continue
            try:
                parsed = json.loads(line)
            except json.JSONDecodeError:
                malformed_lines += 1
                final_line_is_json = False
                continue
            if not isinstance(parsed, dict):
                malformed_lines += 1
                final_line_is_json = False
                continue
            events.append(parsed)
            final_line_is_json = True
            if on_event is not None:
                on_event(parsed)
            yield parsed
        proc.wait()
    finally:
        watchdog.cancel()
        stderr_drain.join(timeout=5.0)

    timed_out = killed.is_set()
    duration_ms = int((time.monotonic() - started) * 1000)
    exit_code = proc.returncode if proc.returncode is not None else -1
    if timed_out:
        exit_code = -1  # adapter-side kill, never a unit's own exit

    terminal = events[-1] if (events and final_line_is_json) else None
    ok = (
        not timed_out
        and exit_code == 0
        and terminal is not None
        and terminal.get("ok") is not False
    )
    yield NativeResult(
        ok=ok,
        exit_code=exit_code,
        data=terminal,
        stderr=stderr_parts[0] if stderr_parts else "",
        duration_ms=duration_ms,
        timed_out=timed_out,
        events=events,
    )