rein-aharness/rein_aharness/request_admission.py

"""Request holds inside the worker's existing private spend ledger.

Provision explicitly before use. Only a trusted owner may bind/revoke a route;
the workload gets an opaque token, never the parent policy, ledger or provider
credential. Lease expiry must come from the accepted queue lease, not a prompt.
"""

from __future__ import annotations

import hashlib
import re
import secrets
import uuid
from datetime import UTC, datetime

from rein_aharness.spend_admission import (
    SpendAdmissionError,
    SpendLedger,
    cap_micros,
    converted_micros,
    timestamp,
)


def _token_hash(token: str) -> str:
    if not isinstance(token, str) or not re.fullmatch(r"[A-Za-z0-9_-]{43,100}", token):
        raise SpendAdmissionError("invalid request route")
    return hashlib.sha256(token.encode()).hexdigest()


def _identity(value: str) -> None:
    if not isinstance(value, str) or not re.fullmatch(
        r"[A-Za-z0-9][A-Za-z0-9._:/@+-]{0,199}", value
    ):
        raise SpendAdmissionError("invalid request identity")


class RequestLedger:
    def __init__(self, parent: SpendLedger):
        self.parent = parent

    def initialize(self) -> None:
        """Explicit, additive schema provisioning; existing request tables refuse."""
        with self.parent._db() as db:
            if db.execute("SELECT 1 FROM reservations WHERE state='held'").fetchone():
                raise SpendAdmissionError("provision request tables before dispatch")
            db.execute("""CREATE TABLE request_routes (
                run_id TEXT PRIMARY KEY, token_sha256 TEXT UNIQUE NOT NULL,
                policy_sha256 TEXT NOT NULL, lease_id TEXT NOT NULL,
                expires_at TEXT NOT NULL, revoked INTEGER NOT NULL CHECK(revoked IN (0,1)))""")
            db.execute("""CREATE TABLE request_reservations (
                receipt TEXT PRIMARY KEY, run_id TEXT NOT NULL,
                policy_sha256 TEXT NOT NULL, lease_id TEXT NOT NULL,
                liability_microusd INTEGER NOT NULL CHECK(liability_microusd>0),
                state TEXT NOT NULL CHECK(state IN ('held','charged')),
                observed_microusd INTEGER, created_at TEXT NOT NULL)""")

    def bind_route(
        self,
        run_id: str,
        policy_sha256: str,
        *,
        lease_id: str,
        expires_at: str,
        now: datetime | None = None,
    ) -> str:
        """Trusted owner only; a parent run can never get a replacement route."""
        _identity(run_id)
        _identity(lease_id)
        if not isinstance(policy_sha256, str) or not re.fullmatch(
            r"[0-9a-f]{64}", policy_sha256
        ):
            raise SpendAdmissionError("invalid request policy digest")
        token = secrets.token_urlsafe(32)
        with self.parent._db() as db:
            clock = now or datetime.now(UTC)
            self.parent._clock(db, clock, admission=True)
            if (
                not clock
                < timestamp(expires_at)
                <= timestamp(self.parent.policy.expires_at)
            ):
                raise SpendAdmissionError("route must expire with the admitted lease")
            row = db.execute(
                "SELECT * FROM reservations WHERE run_id=?", (run_id,)
            ).fetchone()
            if (
                row is None
                or row["state"] != "held"
                or row["definition_id"] != self.parent.policy.activity_definition_id
            ):
                raise SpendAdmissionError("active admitted parent required")
            if db.execute("SELECT breached FROM envelope").fetchone()[0]:
                raise SpendAdmissionError("spend envelope breached")
            db.execute(
                "INSERT INTO request_routes VALUES (?, ?, ?, ?, ?, 0)",
                (run_id, _token_hash(token), policy_sha256, lease_id, expires_at),
            )
        return token

    def revoke_route(self, run_id: str) -> None:
        with self.parent._db() as db:
            changed = db.execute(
                "UPDATE request_routes SET revoked=1 WHERE run_id=?", (run_id,)
            )
            if changed.rowcount != 1:
                raise SpendAdmissionError("unknown request route")

    def _active(self, db, row, now):
        self.parent._clock(db, now, admission=True)
        if row is None or row["revoked"] or now >= timestamp(row["expires_at"]):
            raise SpendAdmissionError("request route expired or revoked")
        parent = db.execute(
            "SELECT state FROM reservations WHERE run_id=?", (row["run_id"],)
        ).fetchone()
        if (
            parent is None
            or parent["state"] != "held"
            or db.execute("SELECT breached FROM envelope").fetchone()[0]
        ):
            raise SpendAdmissionError("request parent not active")

    def reserve_request(
        self,
        token: str,
        policy_sha256: str,
        liability_microusd: int,
        *,
        now: datetime | None = None,
    ) -> str:
        token_hash = _token_hash(token)
        if (
            type(liability_microusd) is not int
            or not 0 < liability_microusd <= 1_000_000_000_000
        ):
            raise SpendAdmissionError("invalid request liability")
        receipt = str(uuid.uuid4())
        with self.parent._db() as db:
            clock = now or datetime.now(UTC)
            route = db.execute(
                "SELECT * FROM request_routes WHERE token_sha256=?", (token_hash,)
            ).fetchone()
            self._active(db, route, clock)
            if route["policy_sha256"] != policy_sha256:
                raise SpendAdmissionError("request policy not admitted")
            if db.execute(
                "SELECT 1 FROM request_reservations WHERE state='held'"
            ).fetchone():
                raise SpendAdmissionError("request outcome unresolved")
            charged = db.execute(
                "SELECT COALESCE(SUM(liability_microusd),0) FROM request_reservations WHERE run_id=?",
                (route["run_id"],),
            ).fetchone()[0]
            if charged + liability_microusd > cap_micros(
                self.parent.policy.max_liability_usd
            ):
                raise SpendAdmissionError("request exceeds remaining parent capacity")
            db.execute(
                "INSERT INTO request_reservations VALUES (?, ?, ?, ?, ?, 'held', NULL, ?)",
                (
                    receipt,
                    route["run_id"],
                    policy_sha256,
                    route["lease_id"],
                    liability_microusd,
                    clock.isoformat(),
                ),
            )
        return receipt

    def request_active(self, receipt: str, *, now: datetime | None = None) -> bool:
        with self.parent._db() as db:
            request = db.execute(
                "SELECT * FROM request_reservations WHERE receipt=?", (receipt,)
            ).fetchone()
            if request is None or request["state"] != "held":
                return False
            route = db.execute(
                "SELECT * FROM request_routes WHERE run_id=?", (request["run_id"],)
            ).fetchone()
            try:
                self._active(db, route, now or datetime.now(UTC))
            except SpendAdmissionError:
                return False
            return True

    def complete_request(
        self, receipt: str, observed_microusd: int, *, now: datetime | None = None
    ) -> None:
        if (
            type(observed_microusd) is not int
            or not 0 <= observed_microusd <= 100_000_000_000_000
        ):
            raise SpendAdmissionError("request accounting incomplete")
        with self.parent._db() as db:
            self.parent._clock(db, now or datetime.now(UTC), admission=False)
            row = db.execute(
                "SELECT * FROM request_reservations WHERE receipt=?", (receipt,)
            ).fetchone()
            if row is None or row["state"] != "held":
                raise SpendAdmissionError("request completion replay refused")
            breached = observed_microusd > row["liability_microusd"]
            db.execute(
                "UPDATE request_reservations SET state=?, liability_microusd=?, observed_microusd=? WHERE receipt=?",
                (
                    "held" if breached else "charged",
                    max(row["liability_microusd"], observed_microusd),
                    observed_microusd,
                    receipt,
                ),
            )
            if breached:
                db.execute("UPDATE envelope SET breached=1")
                db.execute(
                    "UPDATE request_routes SET revoked=1 WHERE run_id=?",
                    (row["run_id"],),
                )
                total = db.execute(
                    "SELECT SUM(liability_microusd) FROM request_reservations WHERE run_id=?",
                    (row["run_id"],),
                ).fetchone()[0]
                # Format integer micros without float rounding.
                usd = f"{total // 1_000_000}.{total % 1_000_000:06d}"
                liability = converted_micros(usd, self.parent.policy.eur_per_usd)
                db.execute(
                    "UPDATE reservations SET liability=MAX(liability, ?) WHERE run_id=?",
                    (liability, row["run_id"]),
                )
        if breached:
            raise SpendAdmissionError("request liability breached; envelope frozen")

    def status(self):
        with self.parent._db() as db:
            return [
                dict(row)
                for row in db.execute(
                    "SELECT * FROM request_reservations ORDER BY created_at, receipt"
                )
            ]