rein-aharness/rein_aharness/owner_bootstrap.py

"""One-cycle entry point for an explicitly admitted exec-env owner child.

Custody/approval belong to the invoking delivery engine. This module neither
fetches a key nor claims that a configuration file authorizes credential access.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
import signal
import threading

from llm_connect.messages_gate import MessagesPolicy
from rein_aharness.messages_owner import MessagesOwner
from rein_aharness.ops_run_client import ActivityCoreOpsClient, OpsRunConfig
from rein_aharness.spend_admission import SpendAdmissionError, _private_file, worker_spend


class BootstrapRefused(RuntimeError):
    pass


def prepare(path: Path, config: OpsRunConfig) -> tuple[MessagesPolicy, Path, str]:
    """Validate value-free, owner-controlled pins without a key or queue claim."""
    from glas_harness.profiles import ProfileCatalog
    from sandboxer.extensions.runtime import verified_runtime
    try:
        _private_file(path)
        if not path.is_absolute() or path.resolve() != path or path.stat().st_size > 65536:
            raise ValueError

        def unique(pairs):
            result = {}
            for key, value in pairs:
                if key in result:
                    raise ValueError
                result[key] = value
            return result

        data = json.loads(path.read_text(), object_pairs_hook=unique)
        if set(data) != {"version", "authority_ref", "spend_policy_sha256", "messages_policy", "runtime"}:
            raise ValueError
        if data["version"] != "1":
            raise ValueError
        spend = worker_spend(config)
        if spend is None or data["authority_ref"] != spend.policy.authority_ref:
            raise ValueError
        if data["spend_policy_sha256"] != spend.policy.sha256:
            raise ValueError
        if (config.worker_id != spend.policy.worker_id
                or config.execution_project != spend.policy.project
                or path.is_relative_to(Path(spend.policy.target_repo).resolve())):
            raise ValueError
        values = data["messages_policy"]
        if not isinstance(values, dict):
            raise ValueError
        if "allowed_betas" in values:
            if not isinstance(values["allowed_betas"], list):
                raise ValueError
            values["allowed_betas"] = tuple(values["allowed_betas"])
        policy = MessagesPolicy(**values)
        runtime = verified_runtime({"runtime": data["runtime"]})
        if runtime is None or not runtime.is_absolute() or runtime.resolve() != runtime:
            raise ValueError
        for private in (path.parent, spend.path.parent, Path(spend.policy.target_repo)):
            a, b = runtime.resolve(), private.resolve()
            if a.is_relative_to(b) or b.is_relative_to(a):
                raise ValueError
        catalog = ProfileCatalog()
        profile, descriptor = catalog.resolve(spend.policy.profile_ref)
        catalog.require_operational(profile)
        from rein_aharness.spend_admission import digest
        if (digest(profile.model_dump(mode="json")) != spend.policy.profile_sha256
                or digest(descriptor.model_dump(mode="json")) != spend.policy.descriptor_sha256
                or profile.model.model != policy.model or profile.credential_route_refs):
            raise ValueError
        from sandboxer.profiles.loader import load_profile
        sandbox = load_profile(profile.sandbox_profile)
        if sandbox.extension != "ext.bwrap" or sandbox.network.default != "deny" or sandbox.network.egress or sandbox.setup.secret_refs:
            raise ValueError
        spend.preflight()
        # Provisioning is a separate owner action; never silently initialize/reset.
        with spend._db() as db:
            db.execute("SELECT run_id, token_sha256, revoked FROM request_routes LIMIT 0")
            db.execute("SELECT receipt, state FROM request_reservations LIMIT 0")
        return policy, runtime, data["runtime"]["sha256"]
    except Exception:
        raise BootstrapRefused("owner bootstrap configuration refused") from None


def run_once(path: Path, *, check_only: bool = False, report_to_hub: bool = True) -> int:
    # Only this explicit exec-env child entry point consumes the conventional key.
    # Remove it before any subprocess or queue client can inherit it. No OAuth,
    # operator HOME, alternate provider URL or credential-file fallback is used.
    provider_key = None if check_only else os.environ.get("ANTHROPIC_API_KEY")
    alternate_auth = None if check_only else os.environ.get("ANTHROPIC_AUTH_TOKEN")
    for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_BASE_URL"):
        os.environ.pop(name, None)
    stop = threading.Event()
    previous = {}
    config = None
    try:
        config = OpsRunConfig.from_env()
        policy, runtime_path, runtime_sha256 = prepare(path, config)
        if check_only:
            print(json.dumps({"ok": True, "check_only": True, "dispatch_enabled": False,
                              "policy_sha256": policy.sha256, "runtime_sha256": runtime_sha256}))
            return 0
        if (alternate_auth or not isinstance(provider_key, str) or not provider_key
                or len(provider_key) > 8192
                or any(ord(c) < 33 or ord(c) > 126 for c in provider_key)):
            raise BootstrapRefused("explicit exec-env provider delivery required")
        owner = MessagesOwner(policy, provider_key, runtime_path=runtime_path,
                              runtime_sha256=runtime_sha256)

        class OneCycleOwner:
            def activate(self, run, config, spend, profile, cancel):
                if stop.is_set():
                    cancel.cancel("signal")
                return owner.activate(run, config, spend, profile, cancel)

        config.messages_owner = OneCycleOwner()
        config.require_request_admission = True
        config.require_spend_admission = True
        client = ActivityCoreOpsClient(config)
        from rein_aharness.claim_loop import _cancel_active_run, process_one
        from rein_aharness.readiness import run_readiness_checks

        def cancelled(signum, frame):
            stop.set()
            _cancel_active_run("signal")

        for signum in (signal.SIGINT, signal.SIGTERM):
            previous[signum] = signal.signal(signum, cancelled)
        if not run_readiness_checks(client).ok:
            raise BootstrapRefused("owner readiness refused")
        if stop.is_set():
            raise BootstrapRefused("owner stopped before claim")
        result = process_one(client, report_to_hub=report_to_hub)
        ok = not stop.is_set() and (result.empty or (result.claimed and result.ok is True))
        print(json.dumps({"ok": ok, "claimed": result.claimed, "empty": result.empty,
                          "run_id": (result.run_id or "")[:200], "ops_state": result.ops_state}))
        return 0 if ok else 1
    except (BootstrapRefused, SpendAdmissionError):
        print(json.dumps({"ok": False, "code": "owner_bootstrap_refused"}))
        return 2
    except Exception:
        # Provider/queue exceptions and config excerpts cannot enter logs/receipts.
        print(json.dumps({"ok": False, "code": "owner_cycle_failed"}))
        return 1
    finally:
        for signum, handler in previous.items():
            signal.signal(signum, handler)
        if config is not None:
            config.messages_owner = None
        provider_key = None