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.

tests/llm/test_fallback.py

@@ -1015,6 +1015,38 @@ class TestLLMCreatorPassesModelUserId:
         assert captured["model_user_id"] == "owner-alice"
 
 
+    @pytest.mark.parametrize(
+        "model_id, source, expected",
+        [(None, None, False), ("catalog-model", "builtin", False), ("byom-uuid", "user", True)],
+    )
+    def test_byom_flag_follows_the_model_source(self, monkeypatch, model_id, source, expected):
+        from types import SimpleNamespace
+
+        from docsgpt.llm.llm_creator import LLMCreator
+        from docsgpt.llm.providers import PROVIDERS_BY_NAME
+
+        class _LLM:
+            def __init__(self, *args, **kwargs):
+                pass
+
+        monkeypatch.setattr(PROVIDERS_BY_NAME["openai"], "llm_class", _LLM)
+        model = SimpleNamespace(
+            source=source, api_key="own-key", base_url=None, upstream_model_id=None, capabilities=None
+        )
+        registry = SimpleNamespace(get_model=lambda _id, user_id=None: model)
+        monkeypatch.setattr(
+            "docsgpt.core.model_registry.ModelRegistry.get_instance", lambda: registry
+        )
+
+        llm = LLMCreator.create_llm(
+            type="openai", api_key="k", user_api_key=None,
+            decoded_token={"sub": "u1"}, model_id=model_id,
+        )
+
+        assert llm._is_byom is expected
+        assert llm._canonical_model_id == model_id
+
+
 # Tests — responding-provider tracking (cross-provider fallback handler fix)
 
 

tests/storage/db/repositories/test_token_usage.py

@@ -35,6 +35,15 @@ class TestInsert:
         )
         assert total == 30
 
+    def test_cost_defaults_to_zero_and_round_trips(self, pg_conn):
+        repo = _repo(pg_conn)
+        repo.insert(user_id="u-cost", prompt_tokens=1, generated_tokens=1)
+        repo.insert(user_id="u-cost", prompt_tokens=1, generated_tokens=1, cost=0.00012345)
+        costs = pg_conn.execute(
+            text("SELECT cost FROM token_usage WHERE user_id = 'u-cost' ORDER BY id")
+        ).scalars().all()
+        assert [float(c) for c in costs] == [0.0, 0.00012345]
+
 
 class TestReassignApiKey:
     def test_rewrites_and_preserves_rate_limit_window(self, pg_conn):

tests/storage/db/test_migration_0033.py

@@ -0,0 +1,145 @@
+"""Migration round-trip test for 0033_quotas."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+from sqlalchemy import text
+from sqlalchemy.exc import IntegrityError
+
+
+pytestmark = pytest.mark.integration
+
+
+def _alembic_ini() -> Path:
+    return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini"
+
+
+def _run_alembic(url: str, *args: str) -> None:
+    subprocess.check_call(
+        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
+        timeout=60,
+        env={**os.environ, "POSTGRES_URI": url},
+    )
+
+
+def _alembic_heads(url: str) -> list[str]:
+    out = subprocess.check_output(
+        [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), "heads"],
+        timeout=60,
+        env={**os.environ, "POSTGRES_URI": url},
+        text=True,
+    )
+    return [line for line in out.splitlines() if line.strip()]
+
+
+def _alembic_version(conn) -> str:
+    return conn.execute(text("SELECT version_num FROM alembic_version")).scalar()
+
+
+def _column_exists(conn, table: str, column: str) -> bool:
+    row = conn.execute(
+        text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'"
+        ),
+        {"t": table, "c": column},
+    ).fetchone()
+    return row is not None
+
+
+_0033 = "0033_quotas"
+_0032 = "0032_personal_access_tokens"
+
+
+def _table_exists(conn, table: str) -> bool:
+    return conn.execute(text("SELECT to_regclass(:t)"), {"t": f"public.{table}"}).scalar() is not None
+
+
+def _insert_policy(conn, **values) -> None:
+    cols = ", ".join(values)
+    params = ", ".join(f":{k}" for k in values)
+    conn.execute(text(f"INSERT INTO quota_policies ({cols}) VALUES ({params})"), values)
+
+
+class TestMigration0033RoundTrip:
+    def test_single_head(self, pg_engine):
+        url = pg_engine.url.render_as_string(hide_password=False)
+        assert len(_alembic_heads(url)) == 1
+
+    def test_head_has_quota_schema(self, pg_engine):
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) >= _0033
+            assert _table_exists(conn, "quota_policies")
+            assert _column_exists(conn, "token_usage", "cost")
+
+    def test_downgrade_drops_then_upgrade_restores(self, pg_engine):
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0032)
+        with pg_engine.connect() as conn:
+            assert _alembic_version(conn) == _0032
+            assert not _table_exists(conn, "quota_policies")
+            assert not _column_exists(conn, "token_usage", "cost")
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            assert _table_exists(conn, "quota_policies")
+            assert _column_exists(conn, "token_usage", "cost")
+
+    def test_upgrade_tolerates_an_existing_cost_column(self, pg_engine):
+        """A database that already carries ``token_usage.cost`` upgrades cleanly."""
+        url = pg_engine.url.render_as_string(hide_password=False)
+        _run_alembic(url, "downgrade", _0032)
+        with pg_engine.begin() as conn:
+            conn.execute(text("ALTER TABLE token_usage ADD COLUMN cost NUMERIC(12,8) NOT NULL DEFAULT 0"))
+            conn.execute(
+                text(
+                    "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens, cost) "
+                    "VALUES ('u-mig33', 10, 1, 0.5)"
+                )
+            )
+        _run_alembic(url, "upgrade", "head")
+        with pg_engine.connect() as conn:
+            cost = conn.execute(text("SELECT cost FROM token_usage WHERE user_id = 'u-mig33'")).scalar()
+            assert float(cost) == 0.5
+
+
+class TestQuotaPolicyConstraints:
+    def test_one_row_per_subject_and_bucket(self, pg_conn):
+        _insert_policy(pg_conn, scope="instance", token_limit=10)
+        with pytest.raises(IntegrityError):
+            with pg_conn.begin_nested():
+                _insert_policy(pg_conn, scope="instance", token_limit=20)
+        _insert_policy(pg_conn, scope="instance", bucket="agent", token_limit=20)
+
+    @pytest.mark.parametrize(
+        "values",
+        [
+            {"scope": "instance", "subject_id": "u1"},
+            {"scope": "user"},
+            {"scope": "user", "subject_id": "u1", "token_limit": 5, "token_unlimited": True},
+            {"scope": "user", "subject_id": "u1", "cost_limit_usd": 5, "cost_unlimited": True},
+            {"scope": "user", "subject_id": "u1", "token_limit": -1},
+            {"scope": "user", "subject_id": "u1", "bucket": "nope"},
+            {"scope": "org", "subject_id": "u1"},
+        ],
+    )
+    def test_invalid_rows_rejected(self, pg_conn, values):
+        with pytest.raises(IntegrityError):
+            with pg_conn.begin_nested():
+                _insert_policy(pg_conn, **values)
+
+    def test_deleting_a_team_removes_its_policies(self, pg_conn):
+        team_id = pg_conn.execute(
+            text("INSERT INTO teams (name, slug, owner_id) VALUES ('Q', 'q-mig33', 'owner') RETURNING id")
+        ).scalar()
+        _insert_policy(pg_conn, scope="team", subject_id=str(team_id), token_limit=10)
+        _insert_policy(pg_conn, scope="user", subject_id=str(team_id), token_limit=10)
+        pg_conn.execute(text("DELETE FROM teams WHERE id = :id"), {"id": team_id})
+        scopes = pg_conn.execute(
+            text("SELECT scope FROM quota_policies WHERE subject_id = :id"), {"id": str(team_id)}
+        ).scalars().all()
+        assert scopes == ["user"]

tests/test_usage.py

@@ -712,3 +712,56 @@ def test_decorator_omits_cache_bins_when_provider_reports_none(monkeypatch):
     assert row["cache_write_tokens"] is None
     assert llm.emitted[0]["cached_tokens"] is None
     assert llm.emitted[0]["cache_write_tokens"] is None
+
+
+def _persist_with_cost(monkeypatch, llm, cost_fn):
+    from docsgpt.usage import _persist_call_usage
+
+    _install_fake_token_repo(monkeypatch)
+    monkeypatch.setattr("docsgpt.usage.compute_cost_usd", cost_fn)
+    _persist_call_usage(
+        llm, {"prompt_tokens": 1000, "generated_tokens": 10, "cached_tokens": 400}
+    )
+    return _FakeTokenUsageRepo.last_instance.inserted[0]
+
+
+class _CostLLM:
+    decoded_token = {"sub": "user_123"}
+    user_api_key = None
+    agent_id = None
+    _canonical_model_id = "priced-model"
+
+
+@pytest.mark.unit
+def test_persist_prices_the_call_by_canonical_model(monkeypatch):
+    seen = {}
+
+    def cost_fn(model, prompt, generated, cached_tokens=None, cache_write_tokens=None):
+        seen.update(model=model, prompt=prompt, generated=generated, cached=cached_tokens)
+        return 0.0123
+
+    row = _persist_with_cost(monkeypatch, _CostLLM(), cost_fn)
+
+    assert row["cost"] == 0.0123
+    assert seen == {"model": "priced-model", "prompt": 1000, "generated": 10, "cached": 400}
+
+
+@pytest.mark.unit
+def test_persist_records_byom_calls_at_zero_cost(monkeypatch):
+    llm = _CostLLM()
+    llm._is_byom = True
+
+    row = _persist_with_cost(monkeypatch, llm, lambda *a, **k: 9.9)
+
+    assert row["cost"] == 0.0
+    assert row["prompt_tokens"] == 1000
+
+
+@pytest.mark.unit
+def test_persist_keeps_the_row_when_pricing_fails(monkeypatch):
+    def boom(*args, **kwargs):
+        raise RuntimeError("registry unavailable")
+
+    row = _persist_with_cost(monkeypatch, _CostLLM(), boom)
+
+    assert row["cost"] == 0.0