lorebase-keeper 0.2.3__py3-none-any.whl

lorebase_keeper/ratelimit/config.py

@@ -0,0 +1,283 @@
+"""Settings of rate limiting and daily quotas (`keeper.limits.rate`, `keeper.limits.quota`).
+
+Every value has a documented default; deployments override it in the YAML file or with
+`LOREBASE__KEEPER__LIMITS__RATE__...` variables. See `config/lorebase.example.yaml`.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+from typing import Literal
+
+from pydantic import Field
+
+from lorebase_platform.settings import SectionModel
+
+
+class RouteClass(StrEnum):
+    """What a request costs, as far as limits are concerned."""
+
+    READS = "reads"
+    WRITES = "writes"
+    SEARCHES = "searches"
+    CONTEXT = "context"
+    ADMIN = "admin"
+
+
+class BucketSettings(SectionModel):
+    """One token bucket: it refills at `requests_per_minute` and holds at most `burst` tokens."""
+
+    requests_per_minute: float = Field(gt=0, description="Sustained rate: tokens added per minute.")
+    burst: int = Field(ge=1, description="Bucket size: requests allowed back to back.")
+
+
+def _bucket(requests_per_minute: float, burst: int) -> BucketSettings:
+    return BucketSettings(requests_per_minute=requests_per_minute, burst=burst)
+
+
+class ClassOverrides(SectionModel):
+    """Buckets of one Domain that replace the installation-wide ones; unset classes keep them."""
+
+    reads: BucketSettings | None = None
+    writes: BucketSettings | None = None
+    searches: BucketSettings | None = None
+    context: BucketSettings | None = None
+    admin: BucketSettings | None = None
+
+    def get(self, route_class: RouteClass) -> BucketSettings | None:
+        """The override of a class, if any."""
+        found: BucketSettings | None = getattr(self, route_class.value)
+        return found
+
+
+class RouteClassing(SectionModel):
+    """How requests are sorted into classes. Anything not matched is `reads` for GET, HEAD and
+    OPTIONS and `writes` for every other method."""
+
+    admin_tags: tuple[str, ...] = Field(
+        default=("admin",), description="Routes with one of these OpenAPI tags are `admin`."
+    )
+    search_suffixes: tuple[str, ...] = Field(
+        default=("/search",), description="POST routes ending in one of these are `searches`."
+    )
+    context_suffixes: tuple[str, ...] = Field(
+        default=("/context",), description="POST routes ending in one of these are `context`."
+    )
+    read_post_suffixes: tuple[str, ...] = Field(
+        default=("/query", "/traverse", "/neighborhood"),
+        description="POST routes ending in one of these only read: they are `reads`.",
+    )
+
+    def only_reads(self, method: str, route_path: str) -> bool:
+        """Whether a request only reads: GET, HEAD, or a POST search, context or read route
+        (the reads recorded by `keeper.auth.audit.reads`)."""
+        if method in ("GET", "HEAD"):
+            return True
+        return method == "POST" and route_path.endswith(
+            (*self.search_suffixes, *self.context_suffixes, *self.read_post_suffixes)
+        )
+
+
+class AuthFailureSettings(BucketSettings):
+    """Buckets of failed authentications: one per client address and, unless
+    `per_key_prefix_enabled` is off, one per API key prefix.
+
+    Only requests answered 401 after presenting a credential take a token. While either bucket
+    of a request is empty it is answered 429 before its credential is checked.
+    """
+
+    requests_per_minute: float = Field(
+        default=30, gt=0, description="Failed attempts regained per minute."
+    )
+    burst: int = Field(default=30, ge=1, description="Failed attempts allowed back to back.")
+    ipv6_prefix_length: int = Field(
+        default=64,
+        ge=1,
+        le=128,
+        description="IPv6 clients share the bucket of their network of this length, since one "
+        "client usually holds a whole /64.",
+    )
+    per_key_prefix_enabled: bool = Field(
+        default=True,
+        description="Keep a bucket per presented API key prefix. Prefixes are not secret, so "
+        "whoever knows one can keep that key throttled; turn this off to count failures by "
+        "client address only (the per-address bucket is always on).",
+    )
+
+
+class ThrottleAuditSettings(SectionModel):
+    """Audit entry for a principal that keeps being throttled."""
+
+    threshold: int = Field(
+        default=20,
+        ge=0,
+        description="Refusals within `window_s` that write one audit entry (`rate_limit."
+        "throttled`); the count then starts again. 0 disables.",
+    )
+    window_s: float = Field(default=300, gt=0, description="Length of the counting window.")
+    max_tracked: int = Field(default=10_000, ge=1, description="Principals counted at once.")
+
+
+class PostgresRateSettings(SectionModel):
+    """The shared backend (`backend: postgres`): buckets in the UNLOGGED `rate_limit_buckets`
+    table of the application database, leased to each Keeper a few tokens at a time."""
+
+    lease_tokens: int = Field(
+        default=10,
+        ge=1,
+        description="Tokens a Keeper takes from a shared bucket per round trip and then serves "
+        "from memory; at most a quarter of the bucket (at least 1). Larger leases mean fewer "
+        "queries; unused leased tokens are lost when the lease expires, which only makes the "
+        "limit stricter. Failed-authentication buckets always take one token at a time.",
+    )
+    lease_ttl_s: float = Field(
+        default=2.0,
+        gt=0,
+        description="Leased tokens not used within this time are dropped, so an idle Keeper "
+        "does not hold a principal's tokens back.",
+    )
+    timeout_s: float = Field(
+        default=0.25,
+        gt=0,
+        description="Time budget of one query to the shared buckets; slower counts as failed.",
+    )
+    on_error: Literal["local", "allow", "deny"] = Field(
+        default="local",
+        description="When the database fails or is slow: `local` enforces the limits in this "
+        "process meanwhile (as backend local), `allow` lets the request through, `deny` "
+        "answers 429. Failed-authentication buckets always fail closed (429).",
+    )
+    error_retry_after_s: int = Field(
+        default=1,
+        ge=1,
+        description="Retry-After of a request refused because the shared buckets could not be "
+        "read (`deny`, and failed-authentication buckets).",
+    )
+    check_cache_s: float = Field(
+        default=0.5,
+        ge=0,
+        description="How long a reading of a failed-authentication bucket serves the checks "
+        "made before each authentication (which take nothing). Failures themselves are always "
+        "written at once; with N Keepers a client may get up to (N-1) x this much of the "
+        "failure rate beyond the limit.",
+    )
+    auth_failures_timeout_s: float = Field(
+        default=1.0,
+        gt=0,
+        description="Time budget of a query of a failed-authentication bucket, waiting for a "
+        "pooled connection included. Past it the bucket fails closed.",
+    )
+    stale_reading_s: float = Field(
+        default=30.0,
+        ge=0,
+        description="While failed-authentication buckets cannot be read, a reading of the same "
+        "bucket at most this old still decides; only clients without one are refused (429).",
+    )
+    prune_interval_s: float = Field(
+        default=60.0,
+        gt=0,
+        description="How often each Keeper deletes rows of buckets that are full again.",
+    )
+    prune_batch: int = Field(default=1000, ge=1, description="Rows deleted per prune.")
+
+
+class RateLimitSettings(SectionModel):
+    """Token buckets per principal and request class (design 13.6)."""
+
+    enabled: bool = Field(default=True, description="Apply rate limits.")
+    backend: Literal["local", "memory", "postgres"] = Field(
+        default="local",
+        description="Where bucket state lives. `local` (`memory` is an alias) is per process: "
+        "with N Keeper processes a principal may use up to N times the limits. `postgres` "
+        "shares the buckets of every Keeper through the application database (see "
+        "`postgres`); use it with more than one Keeper.",
+    )
+    postgres: PostgresRateSettings = Field(default_factory=PostgresRateSettings)
+    max_buckets: int = Field(
+        default=100_000, ge=1, description="Buckets kept in memory; the least recently used go."
+    )
+    exempt_paths: tuple[str, ...] = Field(
+        default=("/healthz", "/readyz"),
+        description="Route paths (without the API prefix) that are never limited or counted.",
+    )
+    reads: BucketSettings = Field(default_factory=lambda: _bucket(1200, 200))
+    writes: BucketSettings = Field(default_factory=lambda: _bucket(300, 60))
+    searches: BucketSettings = Field(default_factory=lambda: _bucket(300, 60))
+    context: BucketSettings = Field(default_factory=lambda: _bucket(120, 30))
+    admin: BucketSettings = Field(default_factory=lambda: _bucket(120, 30))
+    auth_failures: AuthFailureSettings = Field(
+        default_factory=AuthFailureSettings,
+        description="Failed authentications per client address and per API key prefix.",
+    )
+    domains: dict[str, ClassOverrides] = Field(
+        default_factory=dict,
+        description="Per-Domain overrides by Domain id; requests in these Domains use their own "
+        "buckets (not shared with the principal's other Domains).",
+    )
+    routes: RouteClassing = Field(default_factory=RouteClassing)
+    throttle_audit: ThrottleAuditSettings = Field(default_factory=ThrottleAuditSettings)
+
+    def bucket(self, route_class: RouteClass, domain_id: str | None = None) -> BucketSettings:
+        """The bucket for a class, from the Domain's overrides when it has one."""
+        override = self.domains.get(domain_id) if domain_id else None
+        found = override.get(route_class) if override else None
+        if found is not None:
+            return found
+        installation: BucketSettings = getattr(self, route_class.value)
+        return installation
+
+    def has_override(self, route_class: RouteClass, domain_id: str | None) -> bool:
+        """Whether the Domain overrides the bucket of this class."""
+        override = self.domains.get(domain_id) if domain_id else None
+        return override is not None and override.get(route_class) is not None
+
+
+class KeyQuota(SectionModel):
+    """Quota of one API key; a member left out keeps the installation-wide value, `null`
+    removes the limit."""
+
+    requests_per_day: int | None = Field(default=None, ge=1)
+    context_tokens_per_day: int | None = Field(default=None, ge=1)
+
+
+class QuotaSettings(SectionModel):
+    """Daily quotas per API key, counted in UTC days (design 13.6)."""
+
+    enabled: bool = Field(
+        default=True, description="Count the use of API keys. Needs a database; users are exempt."
+    )
+    requests_per_day: int | None = Field(
+        default=100_000, ge=1, description="Requests per key and day; null: unlimited."
+    )
+    context_tokens_per_day: int | None = Field(
+        default=2_000_000,
+        ge=1,
+        description="Estimated tokens of context packs per key and day; null: unlimited. A pack "
+        "that starts under the limit is delivered whole, so a day may end slightly above it.",
+    )
+    flush_interval_s: float = Field(
+        default=5,
+        gt=0,
+        description="How often a process writes the use it has counted. Several processes "
+        "enforce a quota on what they know, so a key may exceed it by the use of one interval.",
+    )
+    keys: dict[str, KeyQuota] = Field(
+        default_factory=dict,
+        description="Overrides by API key id, as listed by the Domain's API key endpoint.",
+    )
+
+    def requests_limit(self, key_id: str) -> int | None:
+        """Requests per day allowed to a key."""
+        return self._limit(key_id, "requests_per_day")
+
+    def tokens_limit(self, key_id: str) -> int | None:
+        """Context-pack tokens per day allowed to a key."""
+        return self._limit(key_id, "context_tokens_per_day")
+
+    def _limit(self, key_id: str, member: str) -> int | None:
+        override = self.keys.get(key_id)
+        if override is not None and member in override.model_fields_set:
+            value: int | None = getattr(override, member)
+            return value
+        default: int | None = getattr(self, member)
+        return default

lorebase_keeper/ratelimit/enforce.py

@@ -0,0 +1,393 @@
+"""Enforcement of rate limits and quotas on requests.
+
+`enforce_limits` is an application-wide dependency: for every route that is not exempt it
+authenticates the caller (refusing clients whose failed authentications used up their
+`auth_failures` buckets before their credential is checked), takes a token from the caller's
+bucket for the request's class and, for API keys, counts the request against the daily quota.
+Refusals are 429 problem details with `Retry-After`. `RateLimitHeadersMiddleware` adds the
+`RateLimit-*` headers of the bucket to every response of a limited route, errors included.
+Context-pack tokens are counted by the route that builds the pack (`check_context_quota`,
+`charge_context_tokens`).
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+from typing import Any, cast
+
+from fastapi import Request, Response
+from opentelemetry.metrics import MeterProvider
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+from starlette.datastructures import MutableHeaders
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from lorebase_keeper.auth.audit import suppress
+from lorebase_keeper.auth.deps import current_principal, get_auth
+from lorebase_keeper.auth.principals import Principal
+from lorebase_keeper.auth.runtime import AuthRuntime
+from lorebase_keeper.auth.sessions import PUBLIC_PATHS
+from lorebase_keeper.client_address import client_address
+from lorebase_keeper.errors import RATE_LIMITED, UNAUTHORIZED, ProblemError, ProblemType
+from lorebase_keeper.ratelimit.bucket import Decision, InMemoryRateStore
+from lorebase_keeper.ratelimit.config import QuotaSettings, RateLimitSettings, RouteClass
+from lorebase_keeper.ratelimit.limiter import AUTH_FAILURES, RateLimiter, ThrottleWatch
+from lorebase_keeper.ratelimit.quotas import QuotaKind, QuotaSnapshot, QuotaTracker
+from lorebase_keeper.ratelimit.shared import PostgresRateStore
+from lorebase_platform.settings import ConfigError
+from lorebase_store.auth import NewAuditEntry
+
+logger = logging.getLogger("lorebase.keeper.limits")
+
+QUOTA_EXCEEDED = ProblemType("quota-exceeded", 429, "Daily quota exceeded")
+STATE_KEY = "rate_limit"
+THROTTLED_ACTION = "rate_limit.throttled"
+LIMIT_HEADER = "RateLimit-Limit"
+REMAINING_HEADER = "RateLimit-Remaining"
+RESET_HEADER = "RateLimit-Reset"
+RETRY_HEADER = "Retry-After"
+EXPOSED_HEADERS = (LIMIT_HEADER, REMAINING_HEADER, RESET_HEADER, RETRY_HEADER)
+
+
+@dataclass
+class LimitRuntime:
+    """The limiting components of one application."""
+
+    rate: RateLimitSettings
+    quota: QuotaSettings
+    limiter: RateLimiter | None
+    quotas: QuotaTracker | None
+    watch: ThrottleWatch | None
+    auth: AuthRuntime
+    stores: tuple[PostgresRateStore, ...] = ()
+
+    def exempt(self, route_path: str) -> bool:
+        """Whether the route (path template without the API prefix) is neither limited nor
+        counted."""
+        return route_path in self.rate.exempt_paths
+
+    async def aclose(self) -> None:
+        """Write the counted use that is still in memory."""
+        if self.quotas is not None:
+            await self.quotas.aclose()
+        for store in self.stores:
+            await store.aclose()
+
+
+def _shared_stores(
+    rate: RateLimitSettings,
+    sessions: async_sessionmaker[AsyncSession],
+    meter_provider: MeterProvider | None,
+) -> tuple[PostgresRateStore, PostgresRateStore]:
+    """The shared stores of principals' buckets and of failed authentications."""
+    principals = PostgresRateStore(
+        sessions,
+        rate.postgres,
+        max_buckets=rate.max_buckets,
+        fallback=InMemoryRateStore(rate.max_buckets),
+        prune=True,
+        label="principals",
+        meter_provider=meter_provider,
+    )
+    failures = PostgresRateStore(
+        sessions,
+        rate.postgres,
+        max_buckets=rate.max_buckets,
+        fail_closed=True,
+        lease_tokens=1,
+        label=AUTH_FAILURES,
+        meter_provider=meter_provider,
+    )
+    return principals, failures
+
+
+def build_limits(
+    rate: RateLimitSettings,
+    quota: QuotaSettings,
+    auth: AuthRuntime,
+    *,
+    sessions: async_sessionmaker[AsyncSession] | None = None,
+    meter_provider: MeterProvider | None = None,
+) -> LimitRuntime | None:
+    """The components for the settings; `None` when neither limit applies. The `postgres`
+    backend needs `sessions` on the application database."""
+    limiter: RateLimiter | None = None
+    stores: tuple[PostgresRateStore, ...] = ()
+    if rate.enabled and rate.backend == "postgres":
+        if sessions is None:
+            raise ConfigError(
+                [
+                    "keeper.limits.rate.backend: postgres needs a database; configure the database "
+                    "section or use backend local"
+                ]
+            )
+        stores = _shared_stores(rate, sessions, meter_provider)
+        limiter = RateLimiter(
+            rate, stores[0], failure_store=stores[1], meter_provider=meter_provider
+        )
+    elif rate.enabled:
+        limiter = RateLimiter(
+            rate, InMemoryRateStore(rate.max_buckets), meter_provider=meter_provider
+        )
+    quotas = (
+        QuotaTracker(quota, auth.store, clock=auth.clock, meter_provider=meter_provider)
+        if quota.enabled and auth.store is not None and auth.settings.uses_api_keys
+        else None
+    )
+    if limiter is None and quotas is None:
+        return None
+    audit = rate.throttle_audit
+    watch = (
+        ThrottleWatch(audit.threshold, audit.window_s, audit.max_tracked)
+        if audit.threshold > 0
+        else None
+    )
+    return LimitRuntime(rate, quota, limiter, quotas, watch, auth, stores)
+
+
+def limits_of(request: Request) -> LimitRuntime | None:
+    """The application's limiting components, if any limit is on."""
+    return cast("LimitRuntime | None", getattr(request.app.state, "limits", None))
+
+
+def principal_label(principal: Principal) -> str:
+    """The key of a principal's buckets: the API key id, else the principal id."""
+    return f"key:{principal.key_id}" if principal.key_id else f"sub:{principal.id}"
+
+
+def route_path(request: Request) -> str:
+    """The route template of the request without the API prefix, whether the router reports
+    it with the prefix or not (as `exempt_paths` and `PUBLIC_PATHS` name routes)."""
+    template = getattr(request.scope.get("route"), "path", None)
+    if not isinstance(template, str):
+        template = str(request.scope["path"])
+    keeper = getattr(request.app.state, "keeper_settings", None)
+    prefix = str(getattr(keeper, "api_prefix", "") or "")
+    if prefix and (template == prefix or template.startswith(prefix + "/")):
+        return template[len(prefix) :] or "/"
+    return template
+
+
+async def _limit_by_address(request: Request, runtime: LimitRuntime, path: str) -> None:
+    """Sign-in routes work without a credential: their buckets are per client address."""
+    if runtime.limiter is None:
+        return
+    address = client_address(request)
+    route_class = runtime.limiter.classify(request.method, path)
+    decision = await runtime.limiter.check(f"addr:{address or 'unknown'}", route_class, None)
+    request.scope.setdefault("state", {})[STATE_KEY] = decision
+    if not decision.allowed:
+        suppress(request)
+        raise ProblemError(
+            RATE_LIMITED,
+            f"Rate limit reached for {route_class.value} requests. "
+            f"Retry in {decision.retry_after_s} s.",
+            extensions={"class": route_class.value, "retry_after_s": decision.retry_after_s},
+            headers={RETRY_HEADER: str(decision.retry_after_s)},
+        )
+
+
+async def enforce_limits(request: Request) -> None:
+    """Application-wide dependency; see the module documentation."""
+    runtime = limits_of(request)
+    if runtime is None:
+        return
+    route = request.scope.get("route")
+    path = route_path(request)
+    if runtime.exempt(path):
+        return
+    if path in PUBLIC_PATHS:
+        await _limit_by_address(request, runtime, path)
+        return
+    principal = await _authenticate(request, runtime)
+    if runtime.limiter is not None:
+        route_class = runtime.limiter.classify(
+            request.method, path, getattr(route, "tags", None) or ()
+        )
+        domain_id = request.path_params.get("domain_id")
+        decision = await runtime.limiter.check(principal_label(principal), route_class, domain_id)
+        request.scope.setdefault("state", {})[STATE_KEY] = decision
+        if not decision.allowed:
+            await _refuse(
+                request,
+                runtime,
+                principal,
+                reason=route_class.value,
+                problem=ProblemError(
+                    RATE_LIMITED,
+                    f"Rate limit reached for {route_class.value} requests. "
+                    f"Retry in {decision.retry_after_s} s.",
+                    extensions={
+                        "class": route_class.value,
+                        "retry_after_s": decision.retry_after_s,
+                    },
+                    headers={RETRY_HEADER: str(decision.retry_after_s)},
+                ),
+            )
+    if runtime.quotas is not None and principal.key_id is not None:
+        refused, snapshot = await runtime.quotas.admit(principal.key_id)
+        if refused is not None:
+            await _refuse(
+                request,
+                runtime,
+                principal,
+                reason=f"quota:{refused}",
+                problem=quota_problem(refused, snapshot, runtime.auth),
+            )
+
+
+async def _authenticate(request: Request, runtime: LimitRuntime) -> Principal:
+    """Authenticate the request, unless its client address or API key prefix has used up its
+    failed attempts: then 429, before any lookup or signature check. A 401 takes one failed
+    attempt from both buckets; success takes nothing."""
+    auth = get_auth(request)
+    limiter = runtime.limiter
+    if limiter is None or not auth.authenticator.presents_credentials(request.headers):
+        return await current_principal(request, auth)
+    address = client_address(request)
+    prefix = auth.authenticator.key_prefix(request.headers)
+    keys = limiter.failure_keys(address, prefix)
+    refused = await limiter.failures_exhausted(keys)
+    if refused is not None:
+        suppress(request)
+        logger.debug(
+            "authentication throttled",
+            extra={"client_address": str(address) if address else None, "key_prefix": prefix},
+        )
+        raise ProblemError(
+            RATE_LIMITED,
+            f"Too many failed authentication attempts. Retry in {refused.retry_after_s} s.",
+            extensions={"class": AUTH_FAILURES, "retry_after_s": refused.retry_after_s},
+            headers={RETRY_HEADER: str(refused.retry_after_s)},
+        )
+    try:
+        return await current_principal(request, auth)
+    except ProblemError as exc:
+        if exc.problem is UNAUTHORIZED:
+            await limiter.record_failure(keys)
+        raise
+
+
+def quota_problem(kind: QuotaKind, snapshot: QuotaSnapshot, auth: AuthRuntime) -> ProblemError:
+    """429 for a used-up daily quota, with the time until it resets."""
+    wait = max(1, math.ceil((snapshot.resets_at - auth.clock()).total_seconds()))
+    limit = snapshot.requests_limit if kind == "requests" else snapshot.tokens_limit
+    return ProblemError(
+        QUOTA_EXCEEDED,
+        f"The daily quota of {limit} {kind.replace('_', ' ')} is used up. "
+        f"It resets at {snapshot.resets_at.isoformat()}.",
+        extensions={
+            "quota": kind,
+            "limit": limit,
+            "resets_at": snapshot.resets_at.isoformat(),
+            "retry_after_s": wait,
+        },
+        headers={RETRY_HEADER: str(wait)},
+    )
+
+
+async def _refuse(
+    request: Request,
+    runtime: LimitRuntime,
+    principal: Principal,
+    *,
+    reason: str,
+    problem: ProblemError,
+) -> None:
+    """Report the refusal for the audit log and raise it."""
+    suppress(request)
+    if runtime.watch is not None and (count := runtime.watch.record(principal_label(principal))):
+        details: dict[str, Any] = {
+            "reason": reason,
+            "refusals": count,
+            "window_s": runtime.rate.throttle_audit.window_s,
+            "method": request.method,
+        }
+        if principal.key_id:
+            details["api_key_id"] = principal.key_id
+        domain_id = request.path_params.get("domain_id")
+        logger.warning("principal throttled repeatedly", extra={"principal": principal.id})
+        writer = runtime.auth.audit_writer
+        if writer is not None:
+            await writer.write(
+                NewAuditEntry(
+                    actor=principal.id,
+                    action=THROTTLED_ACTION,
+                    domain_id=str(domain_id) if domain_id else None,
+                    target=principal.key_id,
+                    details=details,
+                )
+            )
+    raise problem
+
+
+async def check_context_quota(request: Request, principal: Principal) -> None:
+    """429 when the API key has used up its daily context-pack tokens."""
+    runtime = limits_of(request)
+    if runtime is None or runtime.quotas is None or principal.key_id is None:
+        return
+    snapshot = await runtime.quotas.check_tokens(principal.key_id)
+    if snapshot is not None:
+        await _refuse(
+            request,
+            runtime,
+            principal,
+            reason="quota:context_tokens",
+            problem=quota_problem("context_tokens", snapshot, runtime.auth),
+        )
+
+
+async def charge_context_tokens(
+    request: Request, principal: Principal, response: Response, chars_per_token: float
+) -> None:
+    """Count the tokens of a delivered context pack against the API key's daily quota."""
+    runtime = limits_of(request)
+    if runtime is None or runtime.quotas is None or principal.key_id is None:
+        return
+    if response.status_code != 200:
+        return
+    tokens = math.ceil(len(response.body) / chars_per_token)
+    await runtime.quotas.charge_tokens(principal.key_id, tokens)
+
+
+class RateLimitHeadersMiddleware:
+    """Adds `RateLimit-Limit`, `RateLimit-Remaining`, `RateLimit-Reset` (seconds until the bucket
+    is full) to the responses of limited routes."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        async def send_with_headers(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                decision = (scope.get("state") or {}).get(STATE_KEY)
+                if isinstance(decision, Decision):
+                    headers = MutableHeaders(scope=message)
+                    headers[LIMIT_HEADER] = str(decision.limit)
+                    headers[REMAINING_HEADER] = str(decision.remaining)
+                    headers[RESET_HEADER] = str(decision.reset_s)
+                    if not decision.allowed:
+                        headers[RETRY_HEADER] = str(decision.retry_after_s)
+            await send(message)
+
+        await self.app(scope, receive, send_with_headers)
+
+
+__all__ = [
+    "EXPOSED_HEADERS",
+    "QUOTA_EXCEEDED",
+    "THROTTLED_ACTION",
+    "LimitRuntime",
+    "RateLimitHeadersMiddleware",
+    "RouteClass",
+    "build_limits",
+    "charge_context_tokens",
+    "check_context_quota",
+    "enforce_limits",
+    "limits_of",
+]