rein-aharness/rein_aharness/spend_admission.py
tegwick c30806b968
Some checks failed
Governed runtime contract / contract (push) Failing after 31s
Reserve provider requests inside the durable worker envelope
Assistant: codex
Assistant-Model: gpt-5.6-luna
Assistant-Session: 01a07ff8-19d0-7820-b4d0-1353833cb7fc
2026-09-09 21:50:12 +02:00

516 lines
20 KiB
Python

"""Worker-owned, conservative spend reservations; no provider or secret access.
One private SQLite file per envelope, shared by all admitted processes on one
host. Missing state is an error: only the explicit operator init creates it.
"""
from __future__ import annotations
import hashlib
import json
import os
import re
import sqlite3
import stat
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import UTC, datetime
from decimal import ROUND_CEILING, ROUND_FLOOR, Decimal, InvalidOperation, localcontext
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo
class SpendAdmissionError(RuntimeError):
"""A bounded refusal; never contains prompt, credential or provider output."""
def digest(value: Any) -> str:
return hashlib.sha256(
json.dumps(
value, sort_keys=True, separators=(",", ":"), allow_nan=False
).encode()
).hexdigest()
def amount(value: Any, *, positive: bool = True) -> Decimal:
if isinstance(value, bool) or not isinstance(value, (str, int, float)):
raise SpendAdmissionError("invalid monetary amount")
if len(str(value)) > 500:
raise SpendAdmissionError("invalid monetary amount")
try:
result = Decimal(str(value))
except InvalidOperation:
raise SpendAdmissionError("invalid monetary amount") from None
if (
not result.is_finite()
or result < 0
or (positive and result == 0)
or result > 1_000_000
or abs(result.as_tuple().exponent) > 100
):
raise SpendAdmissionError("invalid monetary amount")
return result
def micros(value: Decimal) -> int:
with localcontext() as context:
context.prec = 256
return int((value * 1_000_000).to_integral_value(rounding=ROUND_CEILING))
def cap_micros(value: str) -> int:
with localcontext() as context:
context.prec = 256
return int((amount(value) * 1_000_000).to_integral_value(rounding=ROUND_FLOOR))
def converted_micros(usd: Any, eur_per_usd: str) -> int:
with localcontext() as context:
context.prec = 256
return micros(amount(usd, positive=False) * amount(eur_per_usd))
def timestamp(value: str) -> datetime:
try:
result = datetime.fromisoformat(value)
if result.tzinfo is None:
raise ValueError
return result.astimezone(UTC)
except (TypeError, ValueError):
raise SpendAdmissionError("invalid policy timestamp") from None
@dataclass(frozen=True)
class SpendPolicy:
version: str
envelope_id: str
authority_ref: str
valid_from: str
expires_at: str
timezone: str
worker_id: str
activity_definition_id: str
target_repo: str
project: str
profile_ref: str
profile_sha256: str
descriptor_sha256: str
repository_grant_id: str
max_budget_usd: str
max_liability_usd: str
max_turns: int
eur_per_usd: str
per_run_eur: str
daily_eur: str
total_eur: str
def __post_init__(self) -> None:
if self.version != "1":
raise SpendAdmissionError("unsupported spend policy version")
for name in self.__dataclass_fields__:
value = getattr(self, name)
if name == "max_turns":
if type(value) is not int or not 1 <= value <= 10000:
raise SpendAdmissionError("invalid turn limit")
elif (
not isinstance(value, str)
or not value
or len(value) > 500
or any(ord(c) < 32 for c in value)
):
raise SpendAdmissionError("invalid spend policy field")
for name in ("profile_sha256", "descriptor_sha256"):
if not re.fullmatch(r"[0-9a-f]{64}", getattr(self, name)):
raise SpendAdmissionError("invalid runtime digest")
if not Path(self.target_repo).is_absolute() or "@" not in self.profile_ref:
raise SpendAdmissionError(
"spend policy needs an absolute target and versioned profile"
)
try:
ZoneInfo(self.timezone)
except (ValueError, KeyError):
raise SpendAdmissionError("invalid budget timezone") from None
if timestamp(self.valid_from) >= timestamp(self.expires_at):
raise SpendAdmissionError("empty policy validity interval")
for name in (
"max_budget_usd",
"max_liability_usd",
"eur_per_usd",
"per_run_eur",
"daily_eur",
"total_eur",
):
amount(getattr(self, name))
if amount(self.max_budget_usd) > amount(self.max_liability_usd):
raise SpendAdmissionError("native limit exceeds declared liability")
if self.reservation > cap_micros(self.per_run_eur) or not amount(
self.per_run_eur
) <= amount(self.daily_eur) <= amount(self.total_eur):
raise SpendAdmissionError("liability exceeds operating envelope")
@property
def reservation(self) -> int:
return converted_micros(self.max_liability_usd, self.eur_per_usd)
@property
def sha256(self) -> str:
return digest(self.__dict__)
@classmethod
def load(cls, path: Path) -> SpendPolicy:
_private_file(path)
try:
if path.stat().st_size > 16384:
raise ValueError
def unique(pairs):
result = {}
for key, value in pairs:
if key in result:
raise ValueError
result[key] = value
return result
policy = cls(**json.loads(path.read_text(), object_pairs_hook=unique))
except (OSError, ValueError, TypeError):
raise SpendAdmissionError("invalid spend policy file") from None
if path.resolve().is_relative_to(Path(policy.target_repo).resolve()):
raise SpendAdmissionError("spend policy must be outside target repository")
return policy
def _private_file(path: Path) -> None:
try:
info = path.lstat()
except OSError:
raise SpendAdmissionError("required private spend file unavailable") from None
if (
not stat.S_ISREG(info.st_mode)
or info.st_nlink != 1
or info.st_uid != os.getuid()
or info.st_mode & 0o077
):
raise SpendAdmissionError(
"spend file must be private, regular and worker-owned"
)
class SpendLedger:
def __init__(self, path: Path, policy: SpendPolicy) -> None:
self.path, self.policy = path, policy
parent = path.parent
if not path.is_absolute() or path.resolve().is_relative_to(
Path(policy.target_repo).resolve()
):
raise SpendAdmissionError(
"ledger must be absolute and outside target repository"
)
try:
info = parent.lstat()
except OSError:
raise SpendAdmissionError("private ledger directory unavailable") from None
if (
not stat.S_ISDIR(info.st_mode)
or info.st_uid != os.getuid()
or info.st_mode & 0o077
):
raise SpendAdmissionError(
"ledger directory must be private and worker-owned"
)
def initialize(self) -> None:
"""Explicit provisioning only. Never truncate or recreate existing state."""
try:
fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
os.close(fd)
with self._db(check=False) as db:
db.execute(
"CREATE TABLE envelope (policy TEXT NOT NULL, last_seen TEXT NOT NULL, breached INTEGER NOT NULL CHECK(breached IN (0,1)))"
)
db.execute("""CREATE TABLE reservations (
run_id TEXT PRIMARY KEY, definition_id TEXT NOT NULL,
idempotency_key TEXT NOT NULL, attempt INTEGER NOT NULL,
state TEXT NOT NULL CHECK(state IN ('held','charged')),
liability INTEGER NOT NULL CHECK(liability>0),
start_day TEXT NOT NULL, end_day TEXT,
observed_usd TEXT, receipt TEXT,
UNIQUE(definition_id, idempotency_key),
CHECK((state='held' AND end_day IS NULL) OR
(state='charged' AND end_day IS NOT NULL AND end_day>=start_day)));
""")
db.execute(
"INSERT INTO envelope VALUES (?, ?, 0)",
(self.policy.sha256, self.policy.valid_from),
)
fd = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
try:
os.fsync(fd)
finally:
os.close(fd)
except OSError:
raise SpendAdmissionError("spend ledger initialization refused") from None
@contextmanager
def _db(self, *, check: bool = True) -> Iterator[sqlite3.Connection]:
_private_file(self.path)
db = None
try:
db = sqlite3.connect(
self.path.as_uri() + "?mode=rw",
uri=True,
timeout=10,
isolation_level=None,
)
db.row_factory = sqlite3.Row
db.execute("PRAGMA synchronous=FULL")
db.execute("BEGIN IMMEDIATE")
if check:
rows = db.execute("SELECT * FROM envelope").fetchall()
if len(rows) != 1 or rows[0]["policy"] != self.policy.sha256:
raise SpendAdmissionError("spend ledger policy mismatch")
yield db
db.commit()
except sqlite3.Error:
raise SpendAdmissionError(
"spend ledger unavailable or inconsistent"
) from None
finally:
if db is not None:
db.close() # rolls back any incomplete mutation
def _clock(self, db: sqlite3.Connection, now: datetime, *, admission: bool) -> str:
if now.tzinfo is None:
raise SpendAdmissionError("budget clock must be timezone-aware")
now = now.astimezone(UTC)
previous = timestamp(db.execute("SELECT last_seen FROM envelope").fetchone()[0])
if now < previous:
raise SpendAdmissionError("budget clock moved backwards")
if admission and not timestamp(self.policy.valid_from) <= now < timestamp(
self.policy.expires_at
):
raise SpendAdmissionError("spend policy is not currently valid")
db.execute("UPDATE envelope SET last_seen=?", (now.isoformat(),))
return now.astimezone(ZoneInfo(self.policy.timezone)).date().isoformat()
def _capacity(self, db: sqlite3.Connection, day: str) -> None:
if db.execute("SELECT breached FROM envelope").fetchone()[0]:
raise SpendAdmissionError(
"spend envelope breached; new owner admission required"
)
if db.execute("SELECT 1 FROM reservations WHERE state='held'").fetchone():
raise SpendAdmissionError(
"unresolved spend reservation; reconcile before dispatch"
)
total = db.execute(
"SELECT COALESCE(SUM(liability),0) FROM reservations"
).fetchone()[0]
daily = db.execute(
"SELECT COALESCE(SUM(liability),0) FROM reservations WHERE start_day<=? AND end_day>=?",
(day, day),
).fetchone()[0]
if total + self.policy.reservation > cap_micros(self.policy.total_eur):
raise SpendAdmissionError("total spend capacity exhausted")
if daily + self.policy.reservation > cap_micros(self.policy.daily_eur):
raise SpendAdmissionError("daily spend capacity exhausted")
def preflight(self, *, now: datetime | None = None) -> None:
with self._db() as db:
day = self._clock(db, now or datetime.now(UTC), admission=True)
self._capacity(db, day)
def validate_dispatch(
self, run: Any, config: Any, request: Any, profile: Any, descriptor: Any
) -> None:
p = self.policy
if (
config.worker_id != p.worker_id
or run.claim_owner != p.worker_id
or run.activity_definition_id != p.activity_definition_id
or request.project != p.project
or request.actor != "agt"
or Path(request.repo).resolve() != Path(p.target_repo).resolve()
or run.harness_profile_ref != p.profile_ref
or request.harness_profile_ref != p.profile_ref
or request.request_id != run.id
or not run.repository_grant
or run.repository_grant.grant_id != p.repository_grant_id
or digest(profile.model_dump(mode="json")) != p.profile_sha256
or digest(descriptor.model_dump(mode="json")) != p.descriptor_sha256
):
raise SpendAdmissionError("dispatch does not match admitted spend scope")
if (
profile.rein.id != "rein-aharness"
or amount(profile.limits.max_budget_usd) != amount(p.max_budget_usd)
or profile.limits.max_turns != p.max_turns
):
raise SpendAdmissionError(
"resolved native limits do not match spend policy"
)
def reserve(self, run: Any, *, now: datetime | None = None) -> None:
for value in (run.id, run.activity_definition_id, run.idempotency_key):
if not isinstance(value, str) or not value or len(value) > 500:
raise SpendAdmissionError("invalid spend run identity")
if type(run.attempt) is not int or run.attempt < 1:
raise SpendAdmissionError("invalid spend attempt")
with self._db() as db:
day = self._clock(db, now or datetime.now(UTC), admission=True)
if db.execute(
"SELECT 1 FROM reservations WHERE run_id=? OR (definition_id=? AND idempotency_key=?)",
(run.id, run.activity_definition_id, run.idempotency_key),
).fetchone():
raise SpendAdmissionError(
"run already admitted; workload replay refused"
)
self._capacity(db, day)
db.execute(
"INSERT INTO reservations VALUES (?, ?, ?, ?, 'held', ?, ?, NULL, NULL, NULL)",
(
run.id,
run.activity_definition_id,
run.idempotency_key,
run.attempt,
self.policy.reservation,
day,
),
)
def observe(
self, run_id: str, raw: dict[str, Any], *, now: datetime | None = None
) -> bool:
"""Charge full liability on complete success; never refund self-reported cost."""
evidence = raw.get("evidence", {})
try:
cost = amount(evidence.get("cost_usd"), positive=False)
except SpendAdmissionError:
with self._db() as db:
self._close_requests(db, run_id)
return False # Unknown remains held, even for apparent success.
with self._db() as db:
day = self._clock(db, now or datetime.now(UTC), admission=False)
requests_resolved = self._close_requests(db, run_id)
row = db.execute(
"SELECT * FROM reservations WHERE run_id=?", (run_id,)
).fetchone()
if row is None or row["state"] != "held":
raise SpendAdmissionError("reservation observation conflict")
liability = max(
row["liability"], converted_micros(str(cost), self.policy.eur_per_usd)
)
breached = cost > amount(self.policy.max_liability_usd)
complete = (
raw.get("ok") is True
and evidence.get("outcome") == "succeeded"
and evidence.get("request_id") == run_id
and evidence.get("profile_ref") == self.policy.profile_ref
and evidence.get("session_cleanup") == "succeeded"
and evidence.get("sandbox_destroy") == "succeeded"
and not breached
and requests_resolved
)
db.execute(
"UPDATE reservations SET liability=?, observed_usd=?, state=?, end_day=? WHERE run_id=?",
(
liability,
str(cost),
"charged" if complete else "held",
day if complete else None,
run_id,
),
)
if breached:
db.execute("UPDATE envelope SET breached=1")
return complete
@staticmethod
def _close_requests(
db: sqlite3.Connection, run_id: str, *, reconcile: bool = False
) -> bool:
# Existing ledgers without the opt-in request extension remain compatible.
if not db.execute(
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='request_routes'"
).fetchone():
return True
db.execute("UPDATE request_routes SET revoked=1 WHERE run_id=?", (run_id,))
if reconcile:
# The existing operator reconciliation attests provider termination
# and final accounting. Full child liabilities are retained.
db.execute(
"UPDATE request_reservations SET state='charged' WHERE run_id=?",
(run_id,),
)
return not db.execute(
"SELECT 1 FROM request_reservations WHERE run_id=? AND state='held'",
(run_id,),
).fetchone()
def reconcile(
self, run_id: str, *, cost_usd: str, receipt: str, now: datetime | None = None
) -> None:
"""Operator attests provider termination + final accounting; no retry grant."""
cost = amount(cost_usd, positive=False)
if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._:/@+-]{0,199}", receipt):
raise SpendAdmissionError("bounded reconciliation receipt required")
with self._db() as db:
day = self._clock(db, now or datetime.now(UTC), admission=False)
row = db.execute(
"SELECT * FROM reservations WHERE run_id=?", (run_id,)
).fetchone()
if row is None:
raise SpendAdmissionError("unknown reservation")
if row["state"] != "held":
if row["receipt"] == receipt and row["observed_usd"] == str(cost):
return
raise SpendAdmissionError(
"reconciliation conflicts with settled reservation"
)
if row["observed_usd"] is not None and cost < amount(
row["observed_usd"], positive=False
):
raise SpendAdmissionError(
"reconciliation cannot discard observed liability"
)
self._close_requests(db, run_id, reconcile=True)
liability = max(
row["liability"], converted_micros(str(cost), self.policy.eur_per_usd)
)
db.execute(
"UPDATE reservations SET state='charged', liability=?, observed_usd=?, end_day=?, receipt=? WHERE run_id=?",
(liability, str(cost), day, receipt, run_id),
)
if cost > amount(self.policy.max_liability_usd):
db.execute("UPDATE envelope SET breached=1")
def status(self) -> dict[str, Any]:
with self._db() as db:
return {
"envelope_id": self.policy.envelope_id,
"policy_sha256": self.policy.sha256,
"breached": bool(
db.execute("SELECT breached FROM envelope").fetchone()[0]
),
"reservations": [
dict(row)
for row in db.execute(
"SELECT * FROM reservations ORDER BY start_day, run_id"
)
],
}
def worker_spend(config: Any) -> SpendLedger | None:
if config.spend_policy_path is None and config.spend_ledger_path is None:
if config.require_spend_admission:
raise SpendAdmissionError("required spend admission is not configured")
return None
if not config.spend_policy_path or not config.spend_ledger_path:
raise SpendAdmissionError("both spend policy and ledger must be configured")
policy = SpendPolicy.load(Path(config.spend_policy_path))
if (
config.worker_id != policy.worker_id
or config.execution_project != policy.project
):
raise SpendAdmissionError("worker identity does not match spend policy")
return SpendLedger(Path(config.spend_ledger_path), policy)