feat(quotas): quota_policies table and a per-call cost on token_usage

Migration 0033 adds quota_policies (instance default, team per-member
allowance, user override; a token budget and a USD budget per row) and
token_usage.cost. The column is added IF NOT EXISTS so a database that
already carries it upgrades cleanly.

Every usage row now records the call's USD cost from the model catalog;
bring-your-own models are recorded at $0.

docsgpt/alembic/versions/0033_quotas.py

@@ -0,0 +1,102 @@
+"""0033 quotas — admin-set usage limits and a per-call cost.
+
+``quota_policies`` holds the limits an instance admin sets at three layers:
+the instance default (``subject_id`` NULL), a team's per-member allowance
+(``subject_id`` = ``teams.id``) and a single user's override (``subject_id`` =
+the auth ``sub``). Each row carries a token budget and a USD budget; per
+budget a row either sets a limit (0 blocks), marks it unlimited, or leaves
+both empty to defer to the next layer. ``bucket`` narrows a row to chat without
+an agent (``direct``) or traffic through an agent (``agent``); ``all`` covers both.
+
+``subject_id`` is polymorphic, so there is no FK: an AFTER DELETE trigger on
+``teams`` scrubs a deleted team's rows, and user rows follow the ``user_roles``
+convention of never blocking user deletion.
+
+``token_usage.cost`` is the USD cost of the call at write time (see
+``docsgpt/pricing.py``); 0 for unpriced and bring-your-own models.
+Idempotent both ways.
+
+Revision ID: 0033_quotas
+Revises: 0032_personal_access_tokens
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+
+revision: str = "0033_quotas"
+down_revision: Union[str, None] = "0032_personal_access_tokens"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute(
+        "ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS cost NUMERIC(12,8) NOT NULL DEFAULT 0;"
+    )
+    op.execute(
+        """
+        CREATE TABLE IF NOT EXISTS quota_policies (
+            id              UUID          PRIMARY KEY DEFAULT gen_random_uuid(),
+            scope           TEXT          NOT NULL CHECK (scope IN ('instance', 'team', 'user')),
+            subject_id      TEXT,
+            bucket          TEXT          NOT NULL DEFAULT 'all'
+                                          CHECK (bucket IN ('all', 'direct', 'agent')),
+            token_limit     BIGINT        CHECK (token_limit >= 0),
+            token_unlimited BOOLEAN       NOT NULL DEFAULT false,
+            cost_limit_usd  NUMERIC(12,4) CHECK (cost_limit_usd >= 0),
+            cost_unlimited  BOOLEAN       NOT NULL DEFAULT false,
+            enabled         BOOLEAN       NOT NULL DEFAULT true,
+            note            TEXT,
+            created_by      TEXT,
+            updated_by      TEXT,
+            created_at      TIMESTAMPTZ   NOT NULL DEFAULT now(),
+            updated_at      TIMESTAMPTZ   NOT NULL DEFAULT now(),
+            CONSTRAINT quota_policies_subject_chk
+                CHECK ((scope = 'instance') = (subject_id IS NULL)),
+            CONSTRAINT quota_policies_token_chk
+                CHECK (NOT (token_unlimited AND token_limit IS NOT NULL)),
+            CONSTRAINT quota_policies_cost_chk
+                CHECK (NOT (cost_unlimited AND cost_limit_usd IS NOT NULL))
+        );
+        """
+    )
+    # One row per (layer subject, bucket); the instance row's NULL subject folds to ''.
+    op.execute(
+        "CREATE UNIQUE INDEX IF NOT EXISTS quota_policies_subject_uidx "
+        "ON quota_policies (scope, COALESCE(subject_id, ''), bucket);"
+    )
+    op.execute("DROP TRIGGER IF EXISTS quota_policies_set_updated_at ON quota_policies;")
+    op.execute(
+        """
+        CREATE TRIGGER quota_policies_set_updated_at
+        BEFORE UPDATE ON quota_policies
+        FOR EACH ROW EXECUTE FUNCTION set_updated_at();
+        """
+    )
+    op.execute(
+        """
+        CREATE OR REPLACE FUNCTION cleanup_team_quota_policies() RETURNS trigger AS $$
+        BEGIN
+            DELETE FROM quota_policies WHERE scope = 'team' AND subject_id = OLD.id::text;
+            RETURN OLD;
+        END;
+        $$ LANGUAGE plpgsql;
+        """
+    )
+    op.execute("DROP TRIGGER IF EXISTS teams_cleanup_quota_policies ON teams;")
+    op.execute(
+        """
+        CREATE TRIGGER teams_cleanup_quota_policies
+        AFTER DELETE ON teams
+        FOR EACH ROW EXECUTE FUNCTION cleanup_team_quota_policies();
+        """
+    )
+
+
+def downgrade() -> None:
+    op.execute("DROP TRIGGER IF EXISTS teams_cleanup_quota_policies ON teams;")
+    op.execute("DROP FUNCTION IF EXISTS cleanup_team_quota_policies();")
+    op.execute("DROP TABLE IF EXISTS quota_policies;")
+    op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS cost;")

docsgpt/llm/llm_creator.py

@@ -52,6 +52,7 @@ class LLMCreator:
         base_url = None
         upstream_model_id = model_id
         capabilities = None
+        model = None
         if model_id:
             user_id = model_user_id
             if user_id is None:
@@ -127,4 +128,6 @@ class LLMCreator:
         # llm.model_id is the upstream name (BYOM resolves it above); stamp
         # the canonical id (UUID for BYOM) separately for token_usage.
         llm._canonical_model_id = model_id
+        # Calls to a user's own model are recorded at $0 (see ``docsgpt/usage.py``).
+        llm._is_byom = model is not None and getattr(model, "source", "builtin") == "user"
         return llm

docsgpt/storage/db/models.py

@@ -28,6 +28,7 @@ from sqlalchemy import (
     Index,
     Integer,
     MetaData,
+    Numeric,
     PrimaryKeyConstraint,
     UniqueConstraint,
     Table,
@@ -241,6 +242,9 @@ token_usage_table = Table(
     # cache activity, so hit-rate queries stay honest across providers.
     Column("cached_tokens", Integer),
     Column("cache_write_tokens", Integer),
+    # Added in ``0033_quotas``. USD cost of the call at write time; 0 for
+    # unpriced and bring-your-own models.
+    Column("cost", Numeric(12, 8), nullable=False, server_default="0"),
 )
 
 user_logs_table = Table(
@@ -1121,3 +1125,42 @@ Index(
     personal_access_tokens_table.c.user_id,
     personal_access_tokens_table.c.created_at.desc(),
 )
+
+# --- Usage quotas (migration 0033) ------------------------------------------
+# Admin-set limits at three layers: instance (``subject_id`` NULL), team
+# per-member allowance (``teams.id``) and user override (auth ``sub``). Per
+# budget a row sets a limit, marks it unlimited, or defers to the next layer.
+
+quota_policies_table = Table(
+    "quota_policies",
+    metadata,
+    Column("id", UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()),
+    Column("scope", Text, nullable=False),
+    Column("subject_id", Text),
+    Column("bucket", Text, nullable=False, server_default="all"),
+    Column("token_limit", BigInteger),
+    Column("token_unlimited", Boolean, nullable=False, server_default="false"),
+    Column("cost_limit_usd", Numeric(12, 4)),
+    Column("cost_unlimited", Boolean, nullable=False, server_default="false"),
+    Column("enabled", Boolean, nullable=False, server_default="true"),
+    Column("note", Text),
+    Column("created_by", Text),
+    Column("updated_by", Text),
+    Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
+    Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
+    CheckConstraint("scope IN ('instance', 'team', 'user')", name="quota_policies_scope_check"),
+    CheckConstraint("bucket IN ('all', 'direct', 'agent')", name="quota_policies_bucket_check"),
+    CheckConstraint("token_limit >= 0", name="quota_policies_token_limit_check"),
+    CheckConstraint("cost_limit_usd >= 0", name="quota_policies_cost_limit_check"),
+    CheckConstraint("(scope = 'instance') = (subject_id IS NULL)", name="quota_policies_subject_chk"),
+    CheckConstraint("NOT (token_unlimited AND token_limit IS NOT NULL)", name="quota_policies_token_chk"),
+    CheckConstraint("NOT (cost_unlimited AND cost_limit_usd IS NOT NULL)", name="quota_policies_cost_chk"),
+)
+
+Index(
+    "quota_policies_subject_uidx",
+    quota_policies_table.c.scope,
+    func.coalesce(quota_policies_table.c.subject_id, ""),
+    quota_policies_table.c.bucket,
+    unique=True,
+)

docsgpt/storage/db/repositories/token_usage.py

@@ -37,6 +37,7 @@ class TokenUsageRepository:
         timestamp: Optional[datetime] = None,
         cached_tokens: Optional[int] = None,
         cache_write_tokens: Optional[int] = None,
+        cost: float = 0.0,
     ) -> None:
         # Attribution guard: the ``token_usage_attribution_chk`` CHECK
         # constraint requires at least one of ``user_id`` / ``api_key``
@@ -62,14 +63,14 @@ class TokenUsageRepository:
                 INSERT INTO token_usage (
                     user_id, api_key, agent_id,
                     prompt_tokens, generated_tokens,
-                    cached_tokens, cache_write_tokens,
+                    cached_tokens, cache_write_tokens, cost,
                     source, request_id, model_id, timestamp
                 )
                 VALUES (
                     :user_id, :api_key,
                     CAST(:agent_id AS uuid),
                     :prompt_tokens, :generated_tokens,
-                    :cached_tokens, :cache_write_tokens,
+                    :cached_tokens, :cache_write_tokens, :cost,
                     :source, :request_id, :model_id, COALESCE(:timestamp, now())
                 )
                 """
@@ -82,6 +83,7 @@ class TokenUsageRepository:
                 "generated_tokens": generated_tokens,
                 "cached_tokens": cached_tokens,
                 "cache_write_tokens": cache_write_tokens,
+                "cost": cost,
                 "source": source,
                 "request_id": request_id,
                 "model_id": model_id,

docsgpt/usage.py

@@ -2,6 +2,7 @@ import logging
 import time
 from typing import Any, Dict
 
+from docsgpt.pricing import compute_cost_usd
 from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
 from docsgpt.storage.db.session import db_session
 from docsgpt.utils import num_tokens_from_object_or_list, num_tokens_from_string
@@ -119,6 +120,12 @@ def _persist_call_usage(llm, call_usage):
             },
         )
         return
+    model_id = getattr(llm, "_canonical_model_id", None)
+    # Bring-your-own models run on the user's own provider key: recorded, never priced.
+    if getattr(llm, "_is_byom", False):
+        cost = 0.0
+    else:
+        cost = _call_cost_usd(model_id, call_usage)
     try:
         with db_session() as conn:
             # ``timestamp`` is omitted so Postgres ``server_default
@@ -136,16 +143,32 @@ def _persist_call_usage(llm, call_usage):
                 # "0% cache hits".
                 cached_tokens=call_usage.get("cached_tokens"),
                 cache_write_tokens=call_usage.get("cache_write_tokens"),
+                cost=cost,
                 source=(
                     getattr(llm, "_token_usage_source", None) or "agent_stream"
                 ),
                 request_id=getattr(llm, "_request_id", None),
-                model_id=getattr(llm, "_canonical_model_id", None),
+                model_id=model_id,
             )
     except Exception:
         logger.exception("token_usage persist failed")
 
 
+def _call_cost_usd(model_id, call_usage) -> float:
+    """Price one call; a pricing failure records $0 rather than dropping the row."""
+    try:
+        return compute_cost_usd(
+            model_id,
+            call_usage["prompt_tokens"],
+            call_usage["generated_tokens"],
+            cached_tokens=call_usage.get("cached_tokens"),
+            cache_write_tokens=call_usage.get("cache_write_tokens"),
+        )
+    except Exception:
+        logger.exception("token_usage cost computation failed")
+        return 0.0
+
+
 def _prefer_provider_usage(llm: Any, call_usage: Dict[str, int]) -> Dict[str, int]:
     """Replace estimates with upstream counts when a provider reported them.