Privacy Budget Exhaustion — Detection and Recovery

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Optional


class BudgetExhausted(RuntimeError):
    """Raised when a query's cost would push cumulative ε past the (ε, δ) ceiling.

    Catching this is the fail-closed path: the query is refused and NO noise is
    released, so the accountant's guarantee stays intact.
    """


class RecoveryStrategy(str, Enum):
    FAIL_CLOSED = "fail_closed"   # reject the query outright
    COARSEN = "coarsen"           # lower release resolution → less ε per query
    ROTATE = "rotate"             # new epoch / cohort → fresh window
    CACHED_TILE = "cached_tile"   # serve a pre-aggregated privacy-safe tile


@dataclass
class Accountant:
    """Minimal cumulative-ε accountant. In production, delegate `spent` to an
    RDP/zCDP accountant so composition is counted tightly, not additively."""

    epsilon_max: float
    delta: float = 1e-6
    spent: float = 0.0

    def remaining(self) -> float:
        return self.epsilon_max - self.spent

    def would_exceed(self, cost: float) -> bool:
        # strict '>' — spending exactly to the ceiling is permitted
        return self.spent + cost > self.epsilon_max

    def commit(self, cost: float) -> None:
        """Debit the ledger. MUST be called only after the guard authorises."""
        self.spent += cost


@dataclass
class BudgetExhaustionGuard:
    """Pre-flight budget gate with threshold alerting and recovery strategies.

    `check` is the detection primitive: it fires the threshold alert and then
    raises BEFORE the caller is allowed to draw noise, so a would-be over-spend
    can never leak an under-noised release.
    """

    accountant: Accountant
    alert_fraction: float = 0.80
    coarsen_factor: float = 0.5
    on_threshold: Optional[Callable[[float], None]] = None
    _alerted: bool = field(default=False, repr=False)

    def check(self, cost: float) -> None:
        """Authorise (or veto) a prospective spend of `cost` ε.

        Runs to completion before any noise is drawn. Raises BudgetExhausted
        if committing `cost` would cross epsilon_max.
        """
        if cost < 0:
            raise ValueError("query cost (ε) must be non-negative")
        projected = self.accountant.spent + cost
        fraction = projected / self.accountant.epsilon_max
        # threshold alert fires once per window, at ≥ alert_fraction spent
        if fraction >= self.alert_fraction and not self._alerted:
            self._alerted = True
            if self.on_threshold is not None:
                self.on_threshold(fraction)
        if self.accountant.would_exceed(cost):
            # veto: an unavailable answer is acceptable, an unprotected one is not
            raise BudgetExhausted(
                f"spend {cost:.3f} would exceed ceiling: "
                f"{projected:.3f} > {self.accountant.epsilon_max:.3f}"
            )

    def spend(self, cost: float) -> float:
        """check() then commit(). Returns remaining budget after the debit."""
        self.check(cost)               # gate FIRST — may raise, releasing nothing
        self.accountant.commit(cost)   # only reached if authorised
        return self.accountant.remaining()

    def recover(self, strategy: RecoveryStrategy, cost: float = 0.0) -> float:
        """Apply a recovery strategy; return the ε the caller should now attempt.

        A return of 0.0 means 'no fresh budget is required' (cached or rotated).
        """
        if strategy is RecoveryStrategy.FAIL_CLOSED:
            # refuse and preserve the budget — the default safe action
            raise BudgetExhausted("fail-closed: query rejected, budget preserved")
        if strategy is RecoveryStrategy.COARSEN:
            # coarser grid lowers sensitivity Δf, so the same utility costs less ε
            return cost * self.coarsen_factor
        if strategy is RecoveryStrategy.CACHED_TILE:
            # a pre-aggregated privacy-safe tile was already paid for once
            return 0.0
        if strategy is RecoveryStrategy.ROTATE:
            # new epoch/cohort: open a fresh (ε, δ) window and zero the ledger
            self.accountant.spent = 0.0
            self._alerted = False
            return 0.0
        raise ValueError(f"unknown strategy: {strategy!r}")

def _validate() -> None:
    acc = Accountant(epsilon_max=3.0, delta=1e-6)
    alerts: list[float] = []
    guard = BudgetExhaustionGuard(
        acc, alert_fraction=0.80, coarsen_factor=0.5,
        on_threshold=lambda f: alerts.append(f),
    )
    noise = {"released": False}

    def run_query(cost: float) -> None:
        guard.check(cost)          # 1. gate runs to completion first
        noise["released"] = True   # 2. only reachable when authorised
        acc.commit(cost)

    run_query(1.5)                 # spent 1.5 (50%)
    run_query(1.0)                 # projected 2.5 ≈ 83% → threshold alert fires
    assert alerts and alerts[0] >= 0.80, "threshold alert must fire at ≥ 80% spent"

    # A query that would cross the ceiling is rejected BEFORE any noise is drawn.
    noise["released"] = False
    try:
        run_query(1.0)             # 2.5 + 1.0 = 3.5 > 3.0
    except BudgetExhausted:
        pass
    else:
        raise AssertionError("over-budget query must be rejected")
    assert noise["released"] is False, "no noise may be released on a vetoed query"
    assert acc.spent == 2.5, "a vetoed query must not debit the ledger"

    # Coarsen recovery lowers per-query spend so the retried query now fits.
    cheaper = guard.recover(RecoveryStrategy.COARSEN, cost=1.0)
    assert cheaper < 1.0, "coarsening must reduce ε per query"
    run_query(cheaper)             # 2.5 + 0.5 = 3.0 ≤ 3.0
    assert acc.spent <= acc.epsilon_max

    # Rotate opens a fresh window: the ledger is zeroed for a new epoch.
    guard.recover(RecoveryStrategy.ROTATE)
    assert acc.spent == 0.0 and guard._alerted is False
    print("budget-exhaustion guard invariants hold")


if __name__ == "__main__":
    _validate()