"""HUX-05 run controls: budgets, the pre-side-effect gate and cancellation receipts. The agent hook reports spend per run and asks the gate before every side effect. The gate releases only against an approved, unexpired approval for the same run, capability and argument hash; a ``once`` approval is consumed by its first release (SO-36, SO-37). A stop is not done until a receipt says what actually happened to the run and its side effects (SO-41). """ from __future__ import annotations import re from typing import Any from hux import policy, rules from hux.errors import Invalid from hux.http import Request, Response, Router from hux.identity import Identity from hux.store import TenantStore, check_id BUDGETS = "budgets" RECEIPTS = "receipts" SPEND_KEYS = ("tokens", "tool_calls", "wall_clock_seconds", "delegations", "spend_units", "subagents") LIMIT_OF = { "tokens": "tokens_per_run", "tool_calls": "tool_calls_per_run", "wall_clock_seconds": "wall_clock_seconds", "delegations": "delegations_per_run", "spend_units": "spend_units", "subagents": "subagents_per_run", } RUN_ID_RE = re.compile(r"^[A-Za-z0-9._:-]{1,120}$") def checked_run_id(value: Any) -> str: """Return a canonical run id accepted by both body and path APIs.""" if not isinstance(value, str) or not RUN_ID_RE.fullmatch(value): raise Invalid("run id is malformed or too long") return value def run_id_from(request: Request) -> str: """The run id in the path; the router already bounded its alphabet.""" return checked_run_id(request.params["id"]) # -- budgets ------------------------------------------------------------------- def exhausted(spent: dict[str, int], limits: dict[str, Any]) -> list[str]: """Limit names whose spend has reached them; a zero limit is exhausted immediately.""" return [LIMIT_OF[k] for k in SPEND_KEYS if LIMIT_OF[k] in limits and spent.get(k, 0) >= limits[LIMIT_OF[k]]] def budget_state(store: TenantStore, identity: Identity, run_id: str, conversation_id: str | None = None) -> dict[str, Any]: """Current state for a run, enforced against its conversation's policy epoch aggregate.""" run_id = checked_run_id(run_id) doc_id = f"bud_{policy.run_key(run_id)}" stored = store.get(BUDGETS, doc_id) if store.exists(BUDGETS, doc_id) else {"spent": {}, "_conversation_id": None} known_conversation = stored.get("_conversation_id") or run_conversation(store, run_id) if known_conversation and conversation_id and known_conversation != conversation_id: raise Invalid("run is already bound to another conversation") conversation_id = known_conversation or conversation_id level, scope_id = ("conversation", conversation_id) if conversation_id else ("global", None) effective = policy.effective_policy(store, identity, level, scope_id) limits = {k: v for k, v in effective["budgets"].items() if k != "scope"} epoch = str(effective.get("_budget_epoch") or f"{effective['id']}:{effective['revision']}") same_epoch = stored.get("_budget_epoch") == epoch run_spent = {k: int(stored["spent"].get(k, 0)) if same_epoch else 0 for k in SPEND_KEYS} aggregate = {k: 0 for k in SPEND_KEYS} for record in store.scan(BUDGETS): if record.get("_conversation_id") != conversation_id or record.get("_budget_epoch") != epoch: continue for key in SPEND_KEYS: aggregate[key] += int(record.get("spent", {}).get(key, 0)) return policy.checked({ "schema": "hux.budget_state.v1", "run_id": run_id[:120], "spent": aggregate, "limits": limits, "exhausted": exhausted(aggregate, limits), "_conversation_id": conversation_id, "_id": doc_id, "_run_spent": run_spent, "_budget_epoch": epoch, }) def get_budget(request: Request) -> Response: """``GET /hux/v1/runs/{id}/budget``: spend so far and what is exhausted.""" state = budget_state(request.store, request.identity, run_id_from(request)) request.audit("budgets.read", state["_id"]) return Response(200, policy.public(state)) def post_budget(request: Request) -> Response: """``POST /hux/v1/runs/{id}/budget``: add spend increments; emits budget.exhausted on the crossing.""" policy.require_worker(request.identity, "budget reports") body = policy.body_dict(request) run_id = run_id_from(request) conversation_id = body.get("conversation_id") if conversation_id is not None: conversation_id = check_id(conversation_id) increments = {} for key in SPEND_KEYS: value = body.get(key, 0) if not isinstance(value, int) or isinstance(value, bool) or value < 0: raise Invalid(f"{key} must be a non-negative integer") increments[key] = value with request.store.lock(BUDGETS): before = budget_state(request.store, request.identity, run_id, conversation_id) known = before["exhausted"] if request.store.exists(BUDGETS, before["_id"]) else [] spent = {k: before["_run_spent"][k] + increments[k] for k in SPEND_KEYS} doc = { "id": before["_id"], "run_id": run_id, "spent": spent, "_conversation_id": before["_conversation_id"], "_budget_epoch": before["_budget_epoch"], } request.store.put(BUDGETS, doc) after = budget_state(request.store, request.identity, run_id, conversation_id) request.audit("budgets.write", after["_id"]) newly = [k for k in after["exhausted"] if k not in known] if newly: policy.emit(request.store, request.identity, after["_conversation_id"], "budget.exhausted", f"Budget exhausted: {', '.join(newly)}", run_id=run_id) return Response(200, policy.public(after)) # -- gate ---------------------------------------------------------------------- def hashes_of(approval: dict[str, Any]) -> set[str]: """Argument hashes the approval was requested for (tool_call evidence with a hash).""" return {e["hash"] for e in approval["request"].get("evidence", []) if e.get("kind") == "tool_call" and e.get("hash")} def run_conversation(store: TenantStore, run_id: str) -> str | None: """The conversation a trusted Worker bound to this run in its budget document (F4).""" doc_id = f"bud_{policy.run_key(run_id)}" if store.exists(BUDGETS, doc_id) and store.get(BUDGETS, doc_id).get("_conversation_id"): return store.get(BUDGETS, doc_id)["_conversation_id"] return None def matching_approval(store: TenantStore, run_id: str, capability: str, argument_hash: str, external: bool, conversation_id: str | None) -> tuple[dict[str, Any] | None, str, str | None]: """The approval that releases this side effect, or why none does, plus the conversation the gate settled on. Every release uses an approval record created for this exact run. A session/always grant may let a later run create a new approved record, but the old record itself never crosses run ids. External effects also require the same argument hash; ``once`` names exactly one hash (SO-36, SO-37, SO-39). """ reason = "no approval for this run and capability" for record in store.scan(policy.APPROVALS): record = policy.refresh(store, record) if record["capability"] != capability: continue same_run = record["run_id"] == run_id choice = record.get("decision", {}).get("choice") if not same_run: continue if record["status"] != "approved": reason = f"approval {record['id']} is {record['status']}" continue if policy.parse(record["expires_at"]) <= policy.now(): reason = f"approval {record['id']} has expired" continue if external and not record["request"]["external"]: reason = f"approval {record['id']} was not requested as external" continue if (external or choice == "once") and argument_hash not in hashes_of(record): reason = f"approval {record['id']} was for different arguments" continue if choice == "once" and record.get("_consumed_at"): reason = f"approval {record['id']} was already consumed" continue return record, "released", conversation_id or record["conversation_id"] return None, reason, conversation_id def gate(request: Request) -> Response: """``POST /hux/v1/runs/{id}/gate``: may this side effect proceed right now?""" policy.require_worker(request.identity, "gate checks") body = policy.body_dict(request) run_id = run_id_from(request) capability = body.get("capability") argument_hash = body.get("argument_hash") if capability not in rules.CAPABILITIES: raise Invalid("unknown capability") if not isinstance(argument_hash, str) or not argument_hash.startswith("sha256:"): raise Invalid("argument_hash must be sha256:") external = bool(body.get("external", False)) conversation_id = run_conversation(request.store, run_id) if policy.private_mode_denies(request.store, conversation_id, capability): request.audit("gate.check", f"{run_id}:{capability}", outcome="deny", reason="private_mode") policy.emit(request.store, request.identity, conversation_id, "side_effect.blocked", f"{capability} blocked: private mode", run_id=run_id) return Response(200, {"proceed": False, "reason": "private_mode"}) state = budget_state(request.store, request.identity, run_id, conversation_id) if state["exhausted"]: # F6: an exhausted run releases nothing, whatever was approved request.audit("gate.check", f"{run_id}:{capability}", outcome="deny", reason="budget_exhausted") policy.emit(request.store, request.identity, conversation_id, "budget.exhausted", f"Budget exhausted: {', '.join(state['exhausted'])}", run_id=run_id) return Response(200, {"proceed": False, "reason": "budget_exhausted", "exhausted": state["exhausted"]}) with request.store.lock(policy.APPROVALS): record, reason, conversation_id = matching_approval(request.store, run_id, capability, argument_hash, external, conversation_id) if record is not None and record["decision"]["choice"] == "once": request.store.put(policy.APPROVALS, {**record, "_consumed_at": policy.iso(policy.now())}) if record is None: request.audit("gate.check", f"{run_id}:{capability}", outcome="deny", reason="policy_violation") policy.emit(request.store, request.identity, conversation_id, "side_effect.blocked", f"{capability} blocked: {reason}"[:280], run_id=run_id) return Response(200, {"proceed": False, "reason": reason}) request.audit("gate.check", record["id"]) policy.emit(request.store, request.identity, conversation_id, "side_effect.released", f"{capability} released by approval {record['id']}", run_id=run_id, evidence=[{"kind": "approval", "id": record["id"]}]) return Response(200, {"proceed": True, "approval_id": record["id"], "reason": reason}) # -- stop ---------------------------------------------------------------------- def stop_outcome(body: dict[str, Any]) -> tuple[str, str]: """Return the cancellation outcome proven by the gateway's process registry (F8, SO-41).""" if body.get("already_complete"): return "already_complete", "already_complete" if body.get("process_registry_empty") is not True: return "failed_to_cancel", "process_registry_not_empty" return "cancelled", "cancelled" def stop(request: Request) -> Response: """``POST /hux/v1/runs/{id}/stop``: write the cancellation receipt; a repeat returns it. A ``failed_to_cancel`` receipt is the one non-terminal outcome: a later stop that really cancels or finds the run complete supersedes it with a revision bump; an identical repeat still replays (F8). """ policy.require_worker(request.identity, "stop receipts") body = policy.body_dict(request) run_id = run_id_from(request) receipt_id = f"rcpt_{policy.run_key(run_id)}" side_effects = body.get("side_effects", []) if not isinstance(side_effects, list): raise Invalid("side_effects must be a list") outcome, reason = stop_outcome(body) with request.store.lock(RECEIPTS): existing = request.store.get(RECEIPTS, receipt_id) if request.store.exists(RECEIPTS, receipt_id) else None if existing is not None and (existing["outcome"] != "failed_to_cancel" or outcome == "failed_to_cancel"): request.audit("runs.stop", receipt_id, reason="replayed") return Response(200, policy.public(existing)) stamp = policy.iso(policy.now()) record: dict[str, Any] = { "schema": "hux.cancel_receipt.v1", "id": receipt_id, "run_id": run_id, "requested_by": policy.actor_for(request.identity), "requested_at": existing["requested_at"] if existing else stamp, "acknowledged_at": stamp, "outcome": outcome, "side_effects": side_effects, } if outcome != "failed_to_cancel": record["completed_at"] = stamp conversation_id = body.get("conversation_id", existing.get("conversation_id") if existing else None) if conversation_id is not None: record["conversation_id"] = check_id(conversation_id) stored = request.store.put(RECEIPTS, policy.checked(record), expected_revision=existing["revision"] if existing else None) request.audit("runs.stop", receipt_id, reason=reason if existing is None else f"superseded:{reason}") policy.emit(request.store, request.identity, record.get("conversation_id"), "run.cancelled", f"Run stopped: {outcome} ({reason})", run_id=run_id, evidence=[{"kind": "run", "id": receipt_id}]) return Response(201, policy.public(stored)) def register_routes(router: Router) -> None: """Attach the run-scoped HUX-05 routes; called by ``hux.policy.register``.""" router.add("GET", "/hux/v1/runs/{id}/budget", policy.CARD, "budgets.read", get_budget) router.add("POST", "/hux/v1/runs/{id}/budget", policy.CARD, "budgets.write", post_budget) router.add("POST", "/hux/v1/runs/{id}/gate", policy.CARD, "gate.check", gate) router.add("POST", "/hux/v1/runs/{id}/stop", policy.CARD, "runs.stop", stop)