feat(quotas): limit resolution and the quota service

docsgpt/quotas resolves a user's effective token and cost limits from the
policy rows that apply to them: their own override, then the most generous
allowance among their teams, then the instance default, then any defaults a
registered provider supplies. Each budget resolves on its own, and a team
counts once however many memberships the user holds in it.

QuotaService compares those limits with the user's token_usage totals over
the current QUOTA_PERIOD window (calendar-aligned, UTC, computed at read
time). It fails open, and skips the usage query for unlimited users.

docsgpt/quotas/__init__.py

@@ -0,0 +1,24 @@
+"""Admin-set usage quotas.
+
+Limits live in ``quota_policies`` at three layers (instance default, team
+per-member allowance, user override), each with a token budget and a USD
+budget. ``QuotaService`` resolves a user's effective limits and compares them
+with their ``token_usage`` totals over the current ``QUOTA_PERIOD`` window.
+"""
+
+from docsgpt.quotas.providers import QuotaDefaultsProvider, register_defaults_provider
+from docsgpt.quotas.resolver import ResolvedLimit, ResolvedLimits, resolve_limits
+from docsgpt.quotas.service import BucketStatus, QuotaExceeded, QuotaService
+from docsgpt.quotas.windows import window_bounds
+
+__all__ = [
+    "BucketStatus",
+    "QuotaDefaultsProvider",
+    "QuotaExceeded",
+    "QuotaService",
+    "ResolvedLimit",
+    "ResolvedLimits",
+    "register_defaults_provider",
+    "resolve_limits",
+    "window_bounds",
+]

docsgpt/quotas/providers.py

@@ -0,0 +1,51 @@
+"""Extension point for limits that do not come from ``quota_policies``.
+
+A deployment can register a provider that supplies per-user default policies
+(for example from a subscription plan) and adjusts the quota error payload.
+Defaults sit below every stored layer: a stored instance, team or user row
+with an opinion always wins.
+"""
+
+from __future__ import annotations
+
+from typing import Mapping
+
+
+class QuotaDefaultsProvider:
+    """Base provider: no defaults, error payload unchanged."""
+
+    def default_policies(self, user_id: str) -> list[dict]:
+        """Return default policy rows for ``user_id``.
+
+        Each row uses the ``quota_policies`` field names (``bucket``,
+        ``token_limit``, ``token_unlimited``, ``cost_limit_usd``,
+        ``cost_unlimited``); ``scope`` is set by the caller.
+        """
+        return []
+
+    def error_payload(self, payload: dict, user_id: str) -> dict:
+        """Return the payload sent to a client whose quota is exhausted."""
+        return payload
+
+
+_provider: QuotaDefaultsProvider = QuotaDefaultsProvider()
+
+
+def register_defaults_provider(provider: QuotaDefaultsProvider) -> None:
+    """Replace the process-wide defaults provider."""
+    global _provider
+    _provider = provider
+
+
+def get_defaults_provider() -> QuotaDefaultsProvider:
+    """Return the registered defaults provider."""
+    return _provider
+
+
+def default_rows(user_id: str) -> list[dict]:
+    """Return the provider's defaults for ``user_id`` as ``default``-layer rows."""
+    rows: list[dict] = []
+    for row in get_defaults_provider().default_policies(user_id) or []:
+        if isinstance(row, Mapping):
+            rows.append({"bucket": "all", "enabled": True, **row, "scope": "default", "subject_id": None})
+    return rows

docsgpt/quotas/resolver.py

@@ -0,0 +1,99 @@
+"""Resolve the policy rows that apply to a user into effective limits.
+
+Each budget (tokens, cost) resolves on its own: the user's row wins, then the
+most generous of the user's team rows, then the instance row, then the
+registered defaults. A row with neither a limit nor the unlimited flag for a
+budget has no opinion on it and is skipped.
+
+Teams resolve to the most generous allowance because team membership is not
+controlled by the instance admin: under "most restrictive", any team admin
+could throttle a user by adding them to a low-allowance team.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable, Mapping, Optional
+
+LAYERS = ("user", "team", "instance", "default")
+
+_FIELDS = {"tokens": ("token_limit", "token_unlimited"), "cost": ("cost_limit_usd", "cost_unlimited")}
+
+
+@dataclass(frozen=True)
+class ResolvedLimit:
+    """One budget's effective limit and the layer it came from.
+
+    ``limit`` is ``None`` when the budget is unlimited. ``source`` is ``None``
+    when no layer had an opinion, otherwise one of ``LAYERS``; ``source_id`` is
+    the team id for a team-sourced limit.
+    """
+
+    limit: Optional[float] = None
+    source: Optional[str] = None
+    source_id: Optional[str] = None
+
+    @property
+    def unlimited(self) -> bool:
+        return self.limit is None
+
+
+@dataclass(frozen=True)
+class ResolvedLimits:
+    """A user's effective token and cost limits for one bucket."""
+
+    tokens: ResolvedLimit
+    cost: ResolvedLimit
+
+    @property
+    def unlimited(self) -> bool:
+        return self.tokens.unlimited and self.cost.unlimited
+
+
+def _opinion(row: Mapping, budget: str) -> Optional[tuple[bool, Optional[float]]]:
+    """Return ``(unlimited, limit)`` for a row's budget, or ``None`` if it defers."""
+    limit_field, unlimited_field = _FIELDS[budget]
+    if row.get(unlimited_field):
+        return True, None
+    value = row.get(limit_field)
+    if value is None:
+        return None
+    return False, float(value)
+
+
+def _resolve_budget(rows: Iterable[Mapping], budget: str) -> ResolvedLimit:
+    by_layer: dict[str, list[tuple[Mapping, tuple[bool, Optional[float]]]]] = {}
+    for row in rows:
+        opinion = _opinion(row, budget)
+        if opinion is not None:
+            by_layer.setdefault(row["scope"], []).append((row, opinion))
+    for layer in LAYERS:
+        candidates = by_layer.get(layer)
+        if not candidates:
+            continue
+        # Most generous first: unlimited, then the larger limit. Only the team
+        # layer can hold more than one candidate. Ties break on subject id so
+        # the reported source is stable.
+        row, (unlimited, limit) = min(
+            candidates,
+            key=lambda c: (not c[1][0], -(c[1][1] or 0.0), str(c[0].get("subject_id") or "")),
+        )
+        source_id = str(row["subject_id"]) if layer == "team" else None
+        return ResolvedLimit(limit=None if unlimited else limit, source=layer, source_id=source_id)
+    return ResolvedLimit()
+
+
+def resolve_limits(rows: Iterable[Mapping], bucket: str = "all") -> ResolvedLimits:
+    """Return the effective limits for ``bucket`` from a user's applicable rows.
+
+    Args:
+        rows: Policy rows that apply to the user (their own, their teams', the
+            instance's and any provider defaults). Disabled rows and rows for
+            other buckets are ignored.
+        bucket: The policy bucket to resolve.
+    """
+    applicable = [r for r in rows if r.get("bucket", "all") == bucket and r.get("enabled", True)]
+    return ResolvedLimits(
+        tokens=_resolve_budget(applicable, "tokens"),
+        cost=_resolve_budget(applicable, "cost"),
+    )

docsgpt/quotas/service.py

@@ -0,0 +1,174 @@
+"""Compare a user's usage with their effective limits."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional
+
+from docsgpt.core.settings import settings
+from docsgpt.quotas.providers import default_rows, get_defaults_provider
+from docsgpt.quotas.resolver import ResolvedLimit, ResolvedLimits, resolve_limits
+from docsgpt.quotas.windows import window_bounds
+from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository
+from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
+from docsgpt.storage.db.session import db_readonly
+
+logger = logging.getLogger(__name__)
+
+REQUEST_BUCKETS = ("direct", "agent")
+
+_UNITS = {"tokens": "tokens", "cost": "USD"}
+
+
+@dataclass(frozen=True)
+class BucketStatus:
+    """A user's limits and usage for one policy bucket in the current window."""
+
+    bucket: str
+    limits: ResolvedLimits
+    tokens_used: int
+    cost_used: float
+    resets_at: datetime
+
+    def exceeded_budget(self) -> Optional[str]:
+        """Return ``tokens`` or ``cost`` when that budget is used up, else ``None``."""
+        if not self.limits.tokens.unlimited and self.tokens_used >= self.limits.tokens.limit:
+            return "tokens"
+        if not self.limits.cost.unlimited and self.cost_used >= self.limits.cost.limit:
+            return "cost"
+        return None
+
+    def to_dict(self) -> dict:
+        """Return the JSON shape shared by the admin and user quota endpoints."""
+
+        def budget(limit: ResolvedLimit, used: float) -> dict:
+            return {
+                "limit": limit.limit,
+                "used": used,
+                "source": limit.source,
+                "source_id": limit.source_id,
+            }
+
+        return {
+            "bucket": self.bucket,
+            "tokens": budget(self.limits.tokens, self.tokens_used),
+            "cost": budget(self.limits.cost, round(self.cost_used, 6)),
+            "resets_at": self.resets_at.isoformat(),
+        }
+
+
+@dataclass(frozen=True)
+class QuotaExceeded:
+    """The exhausted budget that blocks a request."""
+
+    user_id: str
+    bucket: str
+    budget: str
+    usage: float
+    limit: float
+    source: Optional[str]
+    source_id: Optional[str]
+    resets_at: datetime
+
+    @property
+    def retry_after_seconds(self) -> int:
+        now = datetime.now(self.resets_at.tzinfo)
+        return max(int((self.resets_at - now).total_seconds()), 1)
+
+    def to_payload(self) -> dict:
+        """Return the client-facing error body, after the provider's adjustments."""
+        unit = _UNITS[self.budget]
+        if self.budget == "cost":
+            amounts = f"${self.usage:.2f} of ${self.limit:.2f}"
+        else:
+            amounts = f"{int(self.usage):,} of {int(self.limit):,} tokens"
+        payload = {
+            "success": False,
+            "error_code": "quota-exceeded",
+            "message": f"Usage quota reached ({amounts}). It resets at {self.resets_at.isoformat()}.",
+            "limit_scope": "user_quota",
+            "dimension": self.budget,
+            "unit": unit,
+            "usage": round(self.usage, 6) if self.budget == "cost" else int(self.usage),
+            "limit": self.limit if self.budget == "cost" else int(self.limit),
+            "bucket": self.bucket,
+            "source": self.source,
+            "resets_at": self.resets_at.isoformat(),
+        }
+        try:
+            return get_defaults_provider().error_payload(payload, self.user_id) or payload
+        except Exception:
+            logger.exception("quota defaults provider failed to build the error payload")
+            return payload
+
+
+class QuotaService:
+    """Resolve limits and measure usage for the current ``QUOTA_PERIOD`` window."""
+
+    @staticmethod
+    def status(
+        user_id: str,
+        buckets: tuple[str, ...] = ("all",),
+        now: Optional[datetime] = None,
+    ) -> list[BucketStatus]:
+        """Return the user's status for each of ``buckets``.
+
+        Usage is only summed for buckets that carry a limit, so a user with no
+        applicable policy costs one policy lookup and no usage query.
+        """
+        start, resets_at = window_bounds(settings.QUOTA_PERIOD, now)
+        statuses: list[BucketStatus] = []
+        with db_readonly() as conn:
+            rows = QuotaPoliciesRepository(conn).policies_for_user(user_id) + default_rows(user_id)
+            usage_repo = TokenUsageRepository(conn)
+            for bucket in buckets:
+                limits = resolve_limits(rows, bucket)
+                tokens_used, cost_used = (0, 0.0)
+                if not limits.unlimited:
+                    tokens_used, cost_used = usage_repo.usage_totals(user_id=user_id, start=start, bucket=bucket)
+                statuses.append(BucketStatus(bucket, limits, tokens_used, cost_used, resets_at))
+        return statuses
+
+    @classmethod
+    def check(
+        cls, user_id: Optional[str], bucket: str = "direct", now: Optional[datetime] = None
+    ) -> Optional[QuotaExceeded]:
+        """Return why ``user_id`` may not start a request, or ``None`` if they may.
+
+        Both the ``all`` policies and the request's own bucket must have room.
+        The check runs before the request, so the call that crosses a limit
+        completes and the next one is refused. Any failure here allows the
+        request: a quota outage must not take chat down.
+
+        Args:
+            user_id: The billable user. Requests with no user are not limited.
+            bucket: ``direct`` for chat without an agent, ``agent`` for traffic through one.
+            now: Reference instant, for tests.
+        """
+        if not user_id:
+            return None
+        if bucket not in REQUEST_BUCKETS:
+            raise ValueError(f"unknown request bucket: {bucket!r}")
+        try:
+            statuses = cls.status(user_id, ("all", bucket), now)
+        except Exception:
+            logger.exception("quota check failed; allowing the request", extra={"user_id": user_id})
+            return None
+        for status in statuses:
+            budget = status.exceeded_budget()
+            if budget is None:
+                continue
+            limit = status.limits.tokens if budget == "tokens" else status.limits.cost
+            return QuotaExceeded(
+                user_id=user_id,
+                bucket=status.bucket,
+                budget=budget,
+                usage=status.tokens_used if budget == "tokens" else status.cost_used,
+                limit=limit.limit,
+                source=limit.source,
+                source_id=limit.source_id,
+                resets_at=status.resets_at,
+            )
+        return None

docsgpt/quotas/windows.py

@@ -0,0 +1,37 @@
+"""Calendar-aligned UTC quota windows, computed at read time (no reset job)."""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+PERIODS = ("day", "week", "month")
+
+
+def window_bounds(period: str, now: Optional[datetime] = None) -> tuple[datetime, datetime]:
+    """Return ``(start, resets_at)`` of the window containing ``now``.
+
+    Args:
+        period: ``day`` (from 00:00), ``week`` (from Monday) or ``month`` (from the 1st).
+        now: Reference instant; defaults to the current time. Naive values are read as UTC.
+
+    Raises:
+        ValueError: If ``period`` is not one of ``PERIODS``.
+    """
+    if now is None:
+        now = datetime.now(timezone.utc)
+    elif now.tzinfo is None:
+        now = now.replace(tzinfo=timezone.utc)
+    else:
+        now = now.astimezone(timezone.utc)
+    midnight = now.replace(hour=0, minute=0, second=0, microsecond=0)
+    if period == "day":
+        return midnight, midnight + timedelta(days=1)
+    if period == "week":
+        start = midnight - timedelta(days=midnight.weekday())
+        return start, start + timedelta(days=7)
+    if period == "month":
+        start = midnight.replace(day=1)
+        next_month = (start + timedelta(days=32)).replace(day=1)
+        return start, next_month
+    raise ValueError(f"unknown quota period: {period!r}")