Design a billing-token rate limiter for AI agents
Budgets in tokens rather than requests, per agent and per tenant, with refunds for failed calls.
- Strategy — Swapping an algorithm at runtime by passing an object.
- Sliding windows and counters — Rate limiters and hit counters as objects with bounded memory.
- Idempotent operations — Making a repeated call harmless inside a single process too.
- Thread-safe collections and atomics — Counters and maps many threads can share.
- Errors, exceptions and validation — Failing loudly at the boundary and never in the middle.
01The design, in the six-step framework
1. Requirements and use cases
A 2026 twist on the rate limiter: AI agents are limited by model tokens, not requests, and the true cost of a call is only known after it finishes. That one fact drives the design: reserve an estimate, then settle.
- An agent reserves an estimated number of tokens before calling a model; the call is refused if a budget would be exceeded.
- When the call ends, the agent settles with the actual tokens used, and the difference is returned or charged.
- A failed call is refunded in full.
- Limits apply at two levels: per agent and per tenant; a call must fit both.
- A settle or refund is idempotent: retrying it does nothing extra.
- A reservation that is never settled expires, so a crashed agent does not hold budget forever.
Out of scope: invoicing, price per model (mentioned as an extension), and sharing across processes. Non-functional constraints: reserve is O(1); the two levels are updated atomically, so a refusal at the tenant level leaves the agent untouched; a long-running call cannot overspend by more than its estimate, plus any overrun the policy allows; and usage over a rolling window is approximate only to one bucket.
2. Core entities and relationships
- TokenLimiter: the facade; owns budgets, reservations and the lock.
- Budget: a rolling-window token allowance for one scope, with used and reserved amounts.
- Reservation: an id, its agent and tenant, the estimate, an expiry, and a status.
- ReservationStatus: HELD, SETTLED, REFUNDED or EXPIRED.
- OverrunPolicy: the strategy for when actual use exceeds the estimate.
- LimitExceeded: the typed error naming which scope refused and when to retry.
The limiter composes budgets and reservations. A Reservation is a small state machine whose terminal states make repeated settles harmless; this is idempotency built into the data rather than bolted on with a set of seen ids. Each Budget is scoped by a string key, so a project level is one more key.
3. Interfaces and key methods
from typing import Protocol
class LimitExceeded(Exception):
"""Carries .scope and .retry_after seconds."""
class OverrunPolicy(Protocol):
def charge(self, estimate: int, actual: int) -> int:
"""Tokens to record when the call finished."""
class TokenLimiter:
def reserve(self, agent: str, tenant: str, estimate: int) -> str:
"""Hold estimate against both scopes atomically; returns a reservation id."""
def settle(self, reservation_id: str, actual: int) -> None:
"""Replace the hold with real usage; a repeat is a no-op."""
def refund(self, reservation_id: str) -> None:
"""Release the hold with no usage; a repeat or a late refund is a no-op."""
def remaining(self, scope: str) -> int:
"""Tokens available now, after holds."""
Invariants: for each budget, held equals the sum of estimates of its HELD reservations; used + held never exceeds the limit at reserve time; a reservation leaves HELD at most once; and any state change touches both scopes inside one critical section.
available(scope) = limit − used in window − held; a reservation passes only if estimate ≤ available for both scopes
4. Design patterns and the choices behind them
Reserve, then settle
The force is an unknown cost at decision time. Charging after the call lets a hundred concurrent long calls each pass a check that sees no spend, and the budget is blown by a hundred times. Charging the estimate up front prevents that. The cost is held budget that may turn out unused, so a large estimate reduces throughput until it settles. This is a two-phase pattern; see distributed transactions for the same shape across services. The rejected option is a per-request limit, which is meaningless when one call can use a thousand times more than another.
Terminal states for idempotency
The force is retries: an agent that times out on settle will call it again, and a double charge is a billing bug. The status field makes this safe, because only a HELD reservation changes anything. A Strategy handles overruns, since the business rule varies: charge the actual amount and let the budget go negative, or cap at the estimate because that is what was promised. A reservation expiring and a late settle arriving is the interesting case: the code records the actual use even after expiry, because the tokens really were consumed, and it does not release the hold a second time.
A rolling window from small buckets
The force is a limit such as one million tokens per hour with bounded memory. Per-minute buckets make usage approximate to a minute and take 60 integers, the approach of sliding window counters. A log of every call would be exact, and it would grow with traffic.
One lock over both scopes
The force is the two-level atomic check. Per-budget locks are possible but need a fixed acquisition order and rollback when the second level refuses. One lock over the short critical section is simpler, and the model call happens outside it, so no lock is held for the seconds a call takes. Typed errors tell the agent which scope refused and when to retry.
5. Core flow and code
Flow for an agent reserving 1,000 and using 700:
reservetakes the lock, expires overdue reservations, and checks both budgets.- Both have room, so
heldrises by 1,000 on each, and the reservation is stored as HELD. The lock is released. - The agent calls the model with no lock held.
settle(rid, 700)releases the 1,000 hold, records 700 in the current bucket of both budgets and marks the reservation SETTLED. A retry finds SETTLED and returns.
import itertools
import threading
from dataclasses import dataclass
from enum import Enum
from typing import Dict, Tuple
class ReservationStatus(Enum):
HELD = 1
SETTLED = 2
REFUNDED = 3
EXPIRED = 4
class LimitExceeded(Exception):
def __init__(self, scope: str, retry_after: float) -> None:
super().__init__(f"{scope} over budget")
self.scope, self.retry_after = scope, retry_after
class ChargeActual:
def charge(self, estimate: int, actual: int) -> int:
return actual
class CapAtEstimate:
def charge(self, estimate: int, actual: int) -> int:
return min(estimate, actual)
class Budget:
def __init__(self, limit: int, window: float, slices: int = 60) -> None:
self.limit, self.window = limit, window
self.width = window / slices
self.slices = slices
self.buckets: Dict[int, int] = {}
self.held = 0
def _used(self, now: float) -> int:
current = int(now // self.width)
for b in [b for b in self.buckets if b <= current - self.slices]:
del self.buckets[b]
return sum(self.buckets.values())
def available(self, now: float) -> int:
return self.limit - self._used(now) - self.held
def spend(self, tokens: int, now: float) -> None:
b = int(now // self.width)
self.buckets[b] = self.buckets.get(b, 0) + tokens
@dataclass
class Reservation:
id: str
scopes: Tuple[str, str]
estimate: int
expires_at: float
status: ReservationStatus = ReservationStatus.HELD
class TokenLimiter:
def __init__(self, clock, agent_limit: int, tenant_limit: int, window: float = 3600.0,
hold_seconds: float = 300.0, policy=None) -> None:
self.clock = clock
self.agent_limit, self.tenant_limit, self.window = agent_limit, tenant_limit, window
self.hold_seconds = hold_seconds
self.policy = policy or ChargeActual()
self.budgets: Dict[str, Budget] = {}
self.reservations: Dict[str, Reservation] = {}
self.ids = itertools.count(1)
self.lock = threading.Lock()
def _budget(self, scope: str) -> Budget:
limit = self.tenant_limit if scope.startswith("tenant:") else self.agent_limit
return self.budgets.setdefault(scope, Budget(limit, self.window))
def reserve(self, agent: str, tenant: str, estimate: int) -> str:
if estimate <= 0:
raise ValueError("estimate must be positive")
with self.lock:
now = self.clock.now()
self._expire(now)
scopes = (f"agent:{agent}", f"tenant:{tenant}")
for scope in scopes: # check both before changing either
if estimate > self._budget(scope).available(now):
raise LimitExceeded(scope, retry_after=self._budget(scope).width)
for scope in scopes:
self._budget(scope).held += estimate
rid = f"r{next(self.ids)}"
self.reservations[rid] = Reservation(rid, scopes, estimate, now + self.hold_seconds)
return rid
def settle(self, reservation_id: str, actual: int) -> None:
with self.lock:
now = self.clock.now()
r = self.reservations[reservation_id]
if r.status in (ReservationStatus.SETTLED, ReservationStatus.REFUNDED):
return
if r.status is ReservationStatus.HELD:
self._release(r)
charged = self.policy.charge(r.estimate, actual)
for scope in r.scopes: # a late settle still records real usage
self._budget(scope).spend(charged, now)
r.status = ReservationStatus.SETTLED
def refund(self, reservation_id: str) -> None:
with self.lock:
r = self.reservations[reservation_id]
if r.status is ReservationStatus.HELD:
self._release(r)
r.status = ReservationStatus.REFUNDED
def remaining(self, scope: str) -> int:
with self.lock:
now = self.clock.now()
self._expire(now)
return self._budget(scope).available(now)
def _release(self, r: Reservation) -> None:
for scope in r.scopes:
self._budget(scope).held -= r.estimate
def _expire(self, now: float) -> None:
for r in self.reservations.values(): # a heap by expiry would avoid this scan
if r.status is ReservationStatus.HELD and r.expires_at <= now:
self._release(r)
r.status = ReservationStatus.EXPIRED
if __name__ == "__main__":
class Clock:
t = 0.0
def now(self) -> float:
return self.t
clock = Clock()
lim = TokenLimiter(clock, agent_limit=1000, tenant_limit=1500)
r1 = lim.reserve("a1", "t", 1000)
assert lim.remaining("agent:a1") == 0 and lim.remaining("tenant:t") == 500
try:
lim.reserve("a2", "t", 600) # agent has room, tenant does not
assert False
except LimitExceeded as e:
assert e.scope == "tenant:t" and lim.remaining("agent:a2") == 1000
lim.settle(r1, 700)
lim.settle(r1, 700) # idempotent
assert lim.remaining("agent:a1") == 300 and lim.remaining("tenant:t") == 800
r2 = lim.reserve("a1", "t", 100)
lim.refund(r2); lim.refund(r2)
assert lim.remaining("agent:a1") == 300
r3 = lim.reserve("a1", "t", 200)
clock.t = 301.0 # crashed agent: hold expires
assert lim.remaining("agent:a1") == 300
lim.settle(r3, 50) # late settle records real usage once
assert lim.remaining("agent:a1") == 250
clock.t = 4000.0 # everything has left the window
assert lim.remaining("agent:a1") == 1000
6. Extensions, concurrency and what interviewers probe
- Estimates. Input tokens are countable up front; output is bounded by
max_tokens. Reserve input plus that bound, and settle returns the rest. - Streaming calls. Reserve in chunks and top up as tokens flow, so a runaway stream is cut when a top-up is refused.
- Different models. Convert tokens to a cost unit with a price table before reserving; the budget stays one number.
- Overrun. Under ChargeActual, a budget can go below zero. That is honest accounting, and later reservations are refused until the window rolls.
- Expiry scan. The code scans for clarity; a heap keyed by
expires_atwith lazy deletion makes it O(log n), as in the TTL store. - Common mistake. Holding the lock across the model call, which serialises all agents for seconds, or checking the second scope after mutating the first, which leaks a hold on refusal.
02Where this question was reported
- PracHub · PracHub Google-tagged (2026-style prompt)