budgets 0.6.0: revocation reaches the meter (revoke, drain, final record, endpoint)

Base protocol §Token Revocation + budgets §Token Scope, drain rule per
AAuth issue #151 (raised from production, not yet in the editor's copy):
- InMemoryMeter.revoke(iss, jti): no new reservations, in-flight complete;
  TokenRevoked(drained) from reserve(); consumed_record withheld until
  drained so the post-revocation record is final.
- make_revocation_endpoint: signed POST {iss, jti}; 200 empty / 404 / 403.
- BudgetMiddleware: AUTH_TOKEN_REVOKED challenge with the final record.
86 tests.

Author: abay94

CHANGELOG.md

@@ -1,5 +1,29 @@
 # Changelog
 
+## 0.6.0
+
+**Budgets — revocation reaches the meter** (base protocol §Token Revocation,
+budgets §Token Scope; drain rule per AAuth issue #151, raised from our
+production and not yet in the editor's copy):
+
+- **`InMemoryMeter.revoke(iss, jti)`**: withdraws the grant so no new request
+  can reserve against the token; requests already in flight complete and are
+  committed as usual. Idempotent; `False` for an unknown `(iss, jti)`.
+- **`TokenRevoked(drained)`**: what `reserve()` returns for a revoked token.
+  `drained` is true once nothing is in flight on it.
+- **Final record only after the drain**: `consumed_record()` withholds a
+  revoked token's record while requests are in flight, so a record the
+  resource issues after the revocation is the token's final figure and the
+  issuer can settle on it (the ordering is checkable by `iat`).
+- **`make_revocation_endpoint(meter, authenticate_ps=…)`**: the base
+  protocol's endpoint — signed `POST {"iss","jti"}`, `200` empty on success or
+  already-invalid, `404` unknown, `403` when the caller is not the issuer.
+- **`BudgetMiddleware`** answers a revoked token with the plain
+  `requirement=auth-token` challenge (`code: AUTH_TOKEN_REVOKED`, no
+  `AAuth-Budget`), the resource token carrying the final record once drained.
+- `revocation_state(key, jti)` for resources that learn of revocation out of
+  band (e.g. from their own registry) and want the same challenge.
+
 ## 0.5.1
 
 **Budgets — the final-record challenge for expired auth tokens** (draft

README.md

@@ -144,6 +144,9 @@ production on [get4agent.com](https://get4agent.com).
 
 - `insufficient-budget` refusals carry **`required`** — the refused request's
   maximum cost — so the agent lowers its bound and retries instead of guessing.
+- **Revocation**: `make_revocation_endpoint(meter, authenticate_ps=…)` is the
+  base protocol's endpoint; a revoked token stops spending at once, in-flight
+  requests complete, and its final record rides on the next challenge.
 - An **expired** auth token gets the base protocol's plain challenge with a
   resource token carrying that token's **final** consumption record — the
   figure its issuer needs to settle the allocation (a record stated at or

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "regent-httpsig"
-version = "0.5.1"
+version = "0.6.0"
 description = "Verify and sign AI agent HTTP traffic in Python — RFC 9421 HTTP Message Signatures: Web Bot Auth (what OpenAI ships) and AAuth."
 readme = "README.md"
 license = "Apache-2.0"

src/regent_httpsig/__init__.py

@@ -9,6 +9,7 @@
     InMemoryMeter,
     InsufficientBudget,
     InvalidBudgetClaim,
+    TokenRevoked,
     UnitMismatch,
 )
 from regent_httpsig.config import HttpsigConfig
@@ -24,13 +25,14 @@
     ResponseSigner,
     UsageQueryError,
     build_usage_response,
+    make_revocation_endpoint,
     make_usage_endpoint,
     parse_usage_request,
     validate_budget_grant,
 )
 from regent_httpsig.verify import WBA_TAG, HttpsigVerifier, VerifiedSignature
 
-__version__ = "0.5.1"
+__version__ = "0.6.0"
 
 __all__ = [
     "DIRECTORY_MEDIA_TYPE",
@@ -40,6 +42,7 @@
     "HttpsigVerifier",
     "InMemoryMeter",
     "InsufficientBudget",
+    "TokenRevoked",
     "InvalidBudgetClaim",
     "NotPublicURL",
     "UnitMismatch",
@@ -51,6 +54,7 @@
     "ResponseSigner",
     "UsageQueryError",
     "build_usage_response",
+    "make_revocation_endpoint",
     "make_usage_endpoint",
     "parse_usage_request",
     "validate_budget_grant",

src/regent_httpsig/budget.py

@@ -43,6 +43,7 @@
     "InsufficientBudget",
     "InvalidBudgetClaim",
     "Reservation",
+    "TokenRevoked",
     "UnitMismatch",
 ]
 
@@ -112,6 +113,17 @@ class InsufficientBudget:
     exhausted: bool
 
 
+@dataclass(frozen=True)
+class TokenRevoked:
+    """Refusal: the presented token was revoked (base protocol §Token
+    Revocation, budgets §Token Scope). The budget ended with the token — no
+    new spend — while requests already in flight complete. ``drained`` says
+    whether those have all settled: only then may the resource state the
+    token's consumption record, which is then FINAL (AAuth issue #151)."""
+
+    drained: bool
+
+
 @dataclass
 class _Pool:
     unit: str
@@ -120,6 +132,7 @@ class _Pool:
     consumed: dict[str, int] = field(default_factory=dict)  # jti -> total committed
     jkt_of: dict[str, str] = field(default_factory=dict)  # jti -> presenting key thumbprint
     reservations: dict[int, tuple[str, int, float]] = field(default_factory=dict)
+    revoked: dict[str, float] = field(default_factory=dict)  # jti -> wall time of revocation
     last_activity: float = 0.0
 
 
@@ -178,6 +191,7 @@ def __init__(self, *, reservation_ttl: float = 120.0,
                  retention_seconds: float = 7200.0,
                  usage_key_retention: float = 86400.0) -> None:
         self._pools: dict[MeterKey, _Pool] = {}
+        self._jti_index: dict[tuple[str, str], MeterKey] = {}  # (iss, jti) -> ledger key
         self._lock = asyncio.Lock()
         self._rids = itertools.count(1)
         self._reservation_ttl = reservation_ttl
@@ -214,9 +228,16 @@ def _purge(self, key: MeterKey, now: float) -> _Pool | None:
         if (not pool.grants and not pool.reservations
                 and now - pool.last_activity > self._retention):
             del self._pools[key]
+            for pair in [p for p, k in self._jti_index.items() if k == key]:
+                del self._jti_index[pair]
             return None
         return pool
 
+    @staticmethod
+    def _drained(pool: _Pool, jti: str) -> bool:
+        """No request is in flight on this token any more."""
+        return not any(j == jti for j, _, _ in pool.reservations.values())
+
     @staticmethod
     def _remaining(pool: _Pool, jti: str) -> int:
         """The presented token's balance: its grant minus what was committed
@@ -252,14 +273,17 @@ async def observe_grant(self, key: MeterKey, jti: str, claim: BudgetClaim,
             pool.last_activity = now
             if jkt:
                 pool.jkt_of.setdefault(jti, jkt)
-            if jti not in pool.grants and wall_delta > 0:
+            self._jti_index.setdefault((key[0], jti), key)
+            if jti not in pool.grants and wall_delta > 0 and jti not in pool.revoked:
                 pool.grants[jti] = (claim.amount, now + wall_delta)
 
     async def reserve(self, key: MeterKey, jti: str,
-                      max_cost: int) -> Reservation | InsufficientBudget:
+                      max_cost: int) -> Reservation | InsufficientBudget | TokenRevoked:
         async with self._lock:
             now = time.monotonic()
             pool = self._purge(key, now)
+            if pool is not None and jti in pool.revoked:
+                return TokenRevoked(drained=self._drained(pool, jti))
             if pool is None or jti not in pool.grants:
                 return InsufficientBudget(remaining=0, exhausted=True)
             remaining = self._remaining(pool, jti)
@@ -302,6 +326,33 @@ async def remaining(self, key: MeterKey, jti: str) -> int:
             pool = self._purge(key, time.monotonic())
             return 0 if pool is None else self._remaining(pool, jti)
 
+    async def revoke(self, iss: str, jti: str) -> bool:
+        """Revoke an auth token by ``(iss, jti)`` — the base protocol's
+        revocation identifier. The grant is withdrawn so no new request can
+        reserve against it; requests already in flight complete and are
+        committed as usual. ``False`` when the pair is unknown here (the
+        endpoint answers 404). Idempotent."""
+        async with self._lock:
+            key = self._jti_index.get((iss, jti))
+            if key is None:
+                return False
+            pool = self._purge(key, time.monotonic())
+            if pool is None:
+                return False
+            pool.revoked.setdefault(jti, time.time())
+            pool.grants.pop(jti, None)
+            pool.last_activity = time.monotonic()
+            return True
+
+    async def revocation_state(self, key: MeterKey, jti: str) -> TokenRevoked | None:
+        """``TokenRevoked`` (with its drain state) if the token was revoked
+        here, else ``None``."""
+        async with self._lock:
+            pool = self._purge(key, time.monotonic())
+            if pool is None or jti not in pool.revoked:
+                return None
+            return TokenRevoked(drained=self._drained(pool, jti))
+
     def _record_usage(self, key: MeterKey, pool: _Pool, jti: str,
                       amount: int) -> None:
         """Post a committed cost to the usage counters (call under lock).
@@ -352,6 +403,10 @@ async def consumed_record(self, key: MeterKey, jti: str) -> dict[str, Any] | Non
             pool = self._purge(key, time.monotonic())
             if pool is None:
                 return None
+            if jti in pool.revoked and not self._drained(pool, jti):
+                # §Token Scope + AAuth #151: a revoked token's record is FINAL,
+                # so it is stated only once nothing is in flight on it.
+                return None
             total = pool.consumed.get(jti, 0)
             return {"jti": jti, "consumed": total} if total > 0 else None
 

src/regent_httpsig/fastapi.py

@@ -41,6 +41,7 @@ async def create_order(sig: VerifiedSignature | None = SignatureDep):
     InvalidBudgetClaim,
     MeterKey,
     Reservation,
+    TokenRevoked,
     UnitMismatch,
 )
 from regent_httpsig.sfv import build_aauth_budget_header, build_aauth_requirement
@@ -229,6 +230,11 @@ async def dispatch(
                                  remaining=0, key=key)
 
         outcome = await self._meter.reserve(key, jti, int(max_cost))
+        if isinstance(outcome, TokenRevoked):
+            # §Token Scope: the budget ended with the token. The challenge is
+            # the base protocol's plain one; the resource token on it carries
+            # the token's FINAL record once nothing is in flight (AAuth #151).
+            return await self._revoked_challenge(key, jti, outcome.drained)
         if isinstance(outcome, InsufficientBudget):
             reason = "budget-exhausted" if outcome.exhausted else "insufficient-budget"
             return await self._refusal_with_token(
@@ -325,6 +331,27 @@ async def _verified(self, request: Request) -> VerifiedSignature | None:
             None if result is not None and result.expired else result)
         return result
 
+    async def _revoked_challenge(self, key: MeterKey, jti: str, drained: bool) -> Response:
+        token: str | None = None
+        if self._resource_token is not None:
+            try:
+                record = await self._meter.consumed_record(key, jti) if drained else None
+                token = await _maybe_await(self._resource_token(key, record))
+            except Exception:  # noqa: BLE001 — the challenge must not fail on the extras
+                logger.warning("resource_token_provider failed", exc_info=True)
+        return JSONResponse(
+            status_code=401,
+            content={
+                "code": "AUTH_TOKEN_REVOKED",
+                "message": "This auth token was revoked by its issuer. Take the resource "
+                           "token in AAuth-Requirement to your PS for a fresh one."
+                           + ("" if drained else " Its final consumption record follows once "
+                              "in-flight requests settle."),
+            },
+            headers={"AAuth-Requirement": build_aauth_requirement(
+                reason=None, resource_token=token)},
+        )
+
     async def _expired_challenge(self, sig: VerifiedSignature) -> Response:
         jti = str(sig.claims.get("jti") or "")
         key: MeterKey = (

src/regent_httpsig/usage.py

@@ -221,3 +221,51 @@ def validate_budget_grant(unit: str, decimals: int,
                     f"(declared in budget_units), got {decimals}")
             return
     raise ValueError(f"unit {unit!r} is not declared in budget_units")
+
+
+def make_revocation_endpoint(
+    meter: Any,
+    *,
+    authenticate_ps: Callable[[Any], Awaitable[str | None]],
+) -> Callable[[Any], Awaitable[Any]]:
+    """The base protocol's revocation endpoint (§Token Revocation) for the
+    budgets meter: a signed ``POST`` identifying an auth token by
+    ``{"iss", "jti"}`` — both REQUIRED, keyed together because a jti is unique
+    only within its issuer. ``200`` when the token was revoked or is already
+    invalid here, ``404`` when the pair is not recognized. The body is
+    deliberately empty on success: the response must not vary with what the
+    recipient holds.
+
+    ``authenticate_ps`` is the same pinned-PS check the usage endpoint uses.
+    The issuer named in the body must be the caller: a PS revokes tokens it
+    issued, never another issuer's. What revocation does to the meter is in
+    :meth:`InMemoryMeter.revoke` — no new spend, in-flight requests complete,
+    and the token's consumption record is withheld until they have (AAuth
+    issue #151: a record issued after the revocation is then final)."""
+    from starlette.responses import JSONResponse, Response
+
+    async def handler(request: Any) -> Response:
+        caller = await authenticate_ps(request)
+        if caller is None:
+            return JSONResponse(status_code=401, content={
+                "code": "PS_AUTH_REQUIRED",
+                "message": "Sign the revocation as the person server that issued the token.",
+            })
+        try:
+            payload = json.loads(await request.body() or b"{}")
+            iss, jti = payload.get("iss"), payload.get("jti")
+            if not isinstance(iss, str) or not isinstance(jti, str) or not iss or not jti:
+                raise ValueError("iss and jti are REQUIRED strings")
+        except (ValueError, json.JSONDecodeError) as exc:
+            return JSONResponse(status_code=400, content={
+                "code": "INVALID_REVOCATION", "message": str(exc)})
+        if iss.rstrip("/") != caller.rstrip("/"):
+            return JSONResponse(status_code=403, content={
+                "code": "NOT_YOUR_TOKEN",
+                "message": "A person server may revoke only tokens it issued."})
+        if not await meter.revoke(iss, jti):
+            return JSONResponse(status_code=404, content={
+                "code": "TOKEN_UNKNOWN", "message": "No such (iss, jti) here."})
+        return Response(status_code=200)
+
+    return handler