From 21f43b2966d93bbc906082673c9e4380cda52c87 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 21 Sep 2026 11:54:44 +0100 Subject: [PATCH] feat(quotas): admin quota API and GET /api/user/quota Admins read and set the instance default, team allowances and user overrides under /api/admin/quotas. A user's endpoint also returns the limits those layers resolve to, the layer each came from and the usage against them. The overview lists catalog models used this period that have no price, since a cost limit cannot see them. Every write is audited. GET /api/user/quota gives a user their own limited buckets, usage and reset time without naming the policies behind them; any valid token may call it. --- docsgpt/api/admin/__init__.py | 1 + docsgpt/api/admin/quotas.py | 276 +++++++++++++++++ docsgpt/api/pat/rules.py | 1 + docsgpt/api/user/me/routes.py | 36 +++ .../storage/db/repositories/token_usage.py | 22 ++ tests/api/test_quota_endpoints.py | 280 ++++++++++++++++++ 6 files changed, 616 insertions(+) create mode 100644 docsgpt/api/admin/quotas.py create mode 100644 tests/api/test_quota_endpoints.py diff --git a/docsgpt/api/admin/__init__.py b/docsgpt/api/admin/__init__.py index 8fa8b444..7bfc3059 100644 --- a/docsgpt/api/admin/__init__.py +++ b/docsgpt/api/admin/__init__.py @@ -1,3 +1,4 @@ from .routes import admin_ns +from . import quotas # noqa: F401 (registers the quota resources on admin_ns) __all__ = ["admin_ns"] diff --git a/docsgpt/api/admin/quotas.py b/docsgpt/api/admin/quotas.py new file mode 100644 index 00000000..8210ba8d --- /dev/null +++ b/docsgpt/api/admin/quotas.py @@ -0,0 +1,276 @@ +"""Admin endpoints for usage quotas (RBAC ``admin`` role required). + +Policies are set at three layers: the instance default, a team's per-member +allowance and a single user's override. Every write is audited to +``auth_events`` with the acting admin recorded. +""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +from flask import jsonify, make_response, request +from flask_restx import Resource + +from docsgpt.api.admin.routes import _actor, admin_ns +from docsgpt.api.user.authz import admin_required +from docsgpt.core.settings import settings +from docsgpt.pricing import is_priced +from docsgpt.quotas.service import REQUEST_BUCKETS, QuotaService +from docsgpt.quotas.windows import window_bounds +from docsgpt.storage.db.base_repository import looks_like_uuid +from docsgpt.storage.db.repositories.auth_events import AuthEventsRepository +from docsgpt.storage.db.repositories.quota_policies import BUCKETS, QuotaPoliciesRepository +from docsgpt.storage.db.repositories.teams import TeamsRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository +from docsgpt.storage.db.repositories.users import UsersRepository +from docsgpt.storage.db.session import db_readonly, db_session + +_MAX_TOKEN_LIMIT = 2**62 +_MAX_COST_LIMIT = 99_999_999.0 +_MAX_NOTE_LENGTH = 500 + + +class _BadPolicy(ValueError): + """The request body does not describe a valid policy.""" + + +def _policy_json(row: dict) -> dict: + cost = row.get("cost_limit_usd") + return { + "scope": row["scope"], + "subject_id": row.get("subject_id"), + "bucket": row["bucket"], + "token_limit": row.get("token_limit"), + "token_unlimited": bool(row.get("token_unlimited")), + "cost_limit_usd": float(cost) if cost is not None else None, + "cost_unlimited": bool(row.get("cost_unlimited")), + "enabled": bool(row.get("enabled", True)), + "note": row.get("note"), + "updated_by": row.get("updated_by"), + "updated_at": row.get("updated_at"), + } + + +def _error(message: str, status: int): + return make_response(jsonify({"success": False, "message": message}), status) + + +def _bucket(value: Any) -> str: + if value is None: + return "all" + if value not in BUCKETS: + raise _BadPolicy(f"bucket must be one of: {', '.join(BUCKETS)}") + return value + + +def _flag(data: dict, key: str, default: bool) -> bool: + value = data.get(key, default) + if not isinstance(value, bool): + raise _BadPolicy(f"{key} must be a boolean") + return value + + +def _parse_policy(data: Any) -> dict: + """Validate a policy body into ``QuotaPoliciesRepository.upsert`` kwargs.""" + if not isinstance(data, dict): + raise _BadPolicy("Body must be a JSON object") + token_limit = data.get("token_limit") + if token_limit is not None: + if isinstance(token_limit, bool) or not isinstance(token_limit, int): + raise _BadPolicy("token_limit must be a whole number or null") + if not 0 <= token_limit <= _MAX_TOKEN_LIMIT: + raise _BadPolicy("token_limit is out of range") + cost_limit = data.get("cost_limit_usd") + if cost_limit is not None: + if isinstance(cost_limit, bool) or not isinstance(cost_limit, (int, float)): + raise _BadPolicy("cost_limit_usd must be a number or null") + if not math.isfinite(cost_limit) or not 0 <= cost_limit <= _MAX_COST_LIMIT: + raise _BadPolicy("cost_limit_usd is out of range") + cost_limit = round(float(cost_limit), 4) + token_unlimited = _flag(data, "token_unlimited", False) + cost_unlimited = _flag(data, "cost_unlimited", False) + if token_unlimited and token_limit is not None: + raise _BadPolicy("Set token_limit or token_unlimited, not both") + if cost_unlimited and cost_limit is not None: + raise _BadPolicy("Set cost_limit_usd or cost_unlimited, not both") + if token_limit is None and cost_limit is None and not token_unlimited and not cost_unlimited: + raise _BadPolicy("Set a limit or mark a budget unlimited; delete the policy to remove it") + note = data.get("note") + if note is not None: + if not isinstance(note, str): + raise _BadPolicy("note must be a string") + note = note.strip()[:_MAX_NOTE_LENGTH] or None + return { + "bucket": _bucket(data.get("bucket")), + "token_limit": token_limit, + "token_unlimited": token_unlimited, + "cost_limit_usd": cost_limit, + "cost_unlimited": cost_unlimited, + "enabled": _flag(data, "enabled", True), + "note": note, + } + + +def _audit(conn, event: str, scope: str, subject_id: Optional[str], detail: dict) -> None: + actor = _actor() + AuthEventsRepository(conn).insert( + # A user policy is filed under that user; the rest under the acting admin. + subject_id if scope == "user" else (actor or "unknown"), + event, + ip=request.remote_addr, + user_agent=request.headers.get("User-Agent"), + metadata={"by": actor, "via": "admin_api", "scope": scope, "subject_id": subject_id, **detail}, + ) + + +def _put_policy(scope: str, subject_id: Optional[str]): + try: + fields = _parse_policy(request.get_json(silent=True)) + except _BadPolicy as exc: + return _error(str(exc), 400) + with db_session() as conn: + row = QuotaPoliciesRepository(conn).upsert( + scope=scope, subject_id=subject_id, actor=_actor(), **fields + ) + _audit(conn, "quota_policy_set", scope, subject_id, fields) + return make_response(jsonify({"success": True, "policy": _policy_json(row)}), 200) + + +def _delete_policy(scope: str, subject_id: Optional[str]): + try: + bucket = _bucket(request.args.get("bucket")) if "bucket" in request.args else None + except _BadPolicy as exc: + return _error(str(exc), 400) + with db_session() as conn: + deleted = QuotaPoliciesRepository(conn).delete(scope, subject_id, bucket) + if deleted: + _audit(conn, "quota_policy_deleted", scope, subject_id, {"bucket": bucket or "*"}) + return make_response(jsonify({"success": True, "deleted": deleted}), 200) + + +def _unpriced_models(conn) -> list[dict]: + """Catalog models used this period that no cost limit can see.""" + start, _ = window_bounds(settings.QUOTA_PERIOD) + return [ + row + for row in TokenUsageRepository(conn).tokens_by_model(start=start) + # BYOM ids are UUIDs; those calls are $0 by design, not by omission. + if not looks_like_uuid(row["model_id"]) and not is_priced(row["model_id"]) + ] + + +@admin_ns.route("/admin/quotas") +class AdminQuotasResource(Resource): + @admin_required + def get(self): + """Every stored policy, grouped by layer, plus the models cost limits cannot see.""" + start, resets_at = window_bounds(settings.QUOTA_PERIOD) + with db_readonly() as conn: + repo = QuotaPoliciesRepository(conn) + teams = {str(t["id"]): t for t in TeamsRepository(conn).list_all()} + team_policies = [] + for row in repo.list_by_scope("team"): + team = teams.get(str(row["subject_id"]), {}) + team_policies.append( + { + **_policy_json(row), + "team_name": team.get("name"), + "team_slug": team.get("slug"), + "member_count": team.get("member_count"), + } + ) + body = { + "success": True, + "period": settings.QUOTA_PERIOD, + "period_start": start.isoformat(), + "resets_at": resets_at.isoformat(), + "instance": [_policy_json(r) for r in repo.list_by_scope("instance")], + "teams": team_policies, + "users": [_policy_json(r) for r in repo.list_by_scope("user")], + "unpriced_models": _unpriced_models(conn), + } + return make_response(jsonify(body), 200) + + +@admin_ns.route("/admin/quotas/instance") +class AdminInstanceQuotaResource(Resource): + @admin_required + def put(self): + """Set the instance default for one bucket.""" + return _put_policy("instance", None) + + @admin_required + def delete(self): + """Remove the instance default for ``?bucket=``, or for every bucket.""" + return _delete_policy("instance", None) + + +@admin_ns.route("/admin/quotas/teams/") +class AdminTeamQuotaResource(Resource): + @admin_required + def get(self, team_id): + """The per-member allowance of one team.""" + if not looks_like_uuid(team_id): + return _error("Team not found", 404) + with db_readonly() as conn: + if TeamsRepository(conn).get(team_id) is None: + return _error("Team not found", 404) + rows = QuotaPoliciesRepository(conn).list_for_subject("team", team_id) + return make_response( + jsonify({"success": True, "policies": [_policy_json(r) for r in rows]}), 200 + ) + + @admin_required + def put(self, team_id): + """Set the allowance each member of the team gets.""" + if not looks_like_uuid(team_id): + return _error("Team not found", 404) + with db_readonly() as conn: + if TeamsRepository(conn).get(team_id) is None: + return _error("Team not found", 404) + return _put_policy("team", team_id) + + @admin_required + def delete(self, team_id): + """Remove the team's allowance for ``?bucket=``, or for every bucket.""" + if not looks_like_uuid(team_id): + return _error("Team not found", 404) + return _delete_policy("team", team_id) + + +@admin_ns.route("/admin/quotas/users/") +class AdminUserQuotaResource(Resource): + @admin_required + def get(self, user_id): + """A user's overrides and the limits and usage they resolve to.""" + with db_readonly() as conn: + if UsersRepository(conn).get(user_id) is None: + return _error("User not found", 404) + rows = QuotaPoliciesRepository(conn).list_for_subject("user", user_id) + statuses = QuotaService.status(user_id, ("all", *REQUEST_BUCKETS)) + return make_response( + jsonify( + { + "success": True, + "period": settings.QUOTA_PERIOD, + "policies": [_policy_json(r) for r in rows], + "effective": [s.to_dict() for s in statuses], + } + ), + 200, + ) + + @admin_required + def put(self, user_id): + """Set one user's override, which beats team allowances and the default.""" + with db_readonly() as conn: + if UsersRepository(conn).get(user_id) is None: + return _error("User not found", 404) + return _put_policy("user", user_id) + + @admin_required + def delete(self, user_id): + """Remove the user's override for ``?bucket=``, or for every bucket.""" + return _delete_policy("user", user_id) diff --git a/docsgpt/api/pat/rules.py b/docsgpt/api/pat/rules.py index 9b69d1a7..227fe11a 100644 --- a/docsgpt/api/pat/rules.py +++ b/docsgpt/api/pat/rules.py @@ -161,6 +161,7 @@ _CHAT = dict( RULES: dict[tuple[str, str], Rule] = { # Identity and public metadata: any valid token. ("/api/user/me", "GET"): _rule(open=True), + ("/api/user/quota", "GET"): _rule(open=True), ("/api/health", "GET"): _rule(open=True), ("/api/config", "GET"): _rule(open=True), # Agents diff --git a/docsgpt/api/user/me/routes.py b/docsgpt/api/user/me/routes.py index 67e31dfb..a6c70d4b 100644 --- a/docsgpt/api/user/me/routes.py +++ b/docsgpt/api/user/me/routes.py @@ -5,6 +5,9 @@ only from ``request.decoded_token`` (already populated and role-resolved by the auth chokepoint in ``app.py``). Auth-mode-agnostic. ``email``/``name``/ ``picture`` are OIDC-only and optional — they are echoed from the token and are never present for ``simple_jwt``/``session_jwt``/no-auth modes. + +``GET /api/user/quota`` returns the caller's usage against the limits an admin +set for them, without naming the policies behind those limits. """ from __future__ import annotations @@ -13,6 +16,8 @@ from flask import jsonify, make_response, request from flask_restx import Namespace, Resource from docsgpt.api.pat.tokens import is_pat +from docsgpt.core.settings import settings +from docsgpt.quotas.service import REQUEST_BUCKETS, QuotaService me_ns = Namespace("me", description="Current user identity and roles", path="/api") @@ -43,3 +48,34 @@ class MeResource(Resource): "resource_filter": decoded_token.get("resource_filter") or {}, } return make_response(jsonify(body), 200) + + +def _own_budget(budget: dict) -> dict: + return {"limit": budget["limit"], "used": budget["used"]} + + +@me_ns.route("/user/quota") +class MyQuotaResource(Resource): + def get(self): + """Return the caller's limited buckets: ``{bucket, tokens, cost, resets_at}`` each.""" + decoded_token = getattr(request, "decoded_token", None) + user_id = decoded_token.get("sub") if decoded_token else None + if not user_id: + return make_response(jsonify({"success": False}), 401) + statuses = QuotaService.status(user_id, ("all", *REQUEST_BUCKETS)) + buckets = [] + for status in statuses: + if status.limits.unlimited: + continue + data = status.to_dict() + buckets.append( + { + "bucket": data["bucket"], + "tokens": _own_budget(data["tokens"]), + "cost": _own_budget(data["cost"]), + "resets_at": data["resets_at"], + } + ) + return make_response( + jsonify({"success": True, "period": settings.QUOTA_PERIOD, "buckets": buckets}), 200 + ) diff --git a/docsgpt/storage/db/repositories/token_usage.py b/docsgpt/storage/db/repositories/token_usage.py index 17e546fe..2c8038fa 100644 --- a/docsgpt/storage/db/repositories/token_usage.py +++ b/docsgpt/storage/db/repositories/token_usage.py @@ -165,6 +165,28 @@ class TokenUsageRepository: ).one() return int(row[0]), float(row[1]) + def tokens_by_model(self, *, start: datetime) -> list[dict]: + """Return ``{model_id, tokens, cost}`` per model since ``start``, busiest first.""" + result = self._conn.execute( + text( + """ + SELECT model_id, + COALESCE(SUM(prompt_tokens + generated_tokens), 0) AS tokens, + COALESCE(SUM(cost), 0) AS cost + FROM token_usage + WHERE timestamp >= :start AND model_id IS NOT NULL + AND source <> ALL(:rollup_sources) + GROUP BY model_id + ORDER BY tokens DESC, model_id + """ + ), + {"start": start, "rollup_sources": list(self.ROLLUP_SOURCES)}, + ) + return [ + {"model_id": row[0], "tokens": int(row[1]), "cost": float(row[2])} + for row in result.fetchall() + ] + # Token usage written outside a user-initiated request (conversation # title generation, history compression, RAG question condensing, # provider fallback). Mirrors the exclusion list in ``count_in_range``. diff --git a/tests/api/test_quota_endpoints.py b/tests/api/test_quota_endpoints.py new file mode 100644 index 00000000..44c424db --- /dev/null +++ b/tests/api/test_quota_endpoints.py @@ -0,0 +1,280 @@ +"""Endpoint tests for the admin quota API and ``GET /api/user/quota``. + +Driven through the real app.py chokepoint against an ephemeral Postgres; only +``handle_auth`` / ``resolve_roles`` are patched. +""" + +from __future__ import annotations + +import json +from contextlib import ExitStack, contextmanager +from unittest.mock import patch + +import pytest +from sqlalchemy import text + +from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository +from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository +from docsgpt.storage.db.repositories.users import UsersRepository + + +@pytest.fixture +def client(): + from docsgpt.app import app as flask_app + + flask_app.config["TESTING"] = True + return flask_app.test_client() + + +@pytest.fixture +def db(pg_conn): + @contextmanager + def _yield(): + yield pg_conn + + with ExitStack() as stack: + for target in ( + "docsgpt.api.admin.quotas.db_readonly", + "docsgpt.api.admin.quotas.db_session", + "docsgpt.quotas.service.db_readonly", + ): + stack.enter_context(patch(target, _yield)) + yield pg_conn + + +@contextmanager +def _as(sub, *roles): + with patch("docsgpt.app.handle_auth", return_value={"sub": sub}), patch( + "docsgpt.app.resolve_roles", return_value=list(roles) or ["user"] + ): + yield + + +def _admin(): + return _as("admin1", "admin", "user") + + +def _body(resp): + return json.loads(resp.data) + + +def _team(conn, slug="q-team", member=None): + team_id = str( + conn.execute( + text("INSERT INTO teams (name, slug, owner_id) VALUES (:n, :s, 'o') RETURNING id"), + {"n": slug, "s": slug}, + ).scalar() + ) + if member: + conn.execute( + text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), :u, 'team_member')"), + {"t": team_id, "u": member}, + ) + return team_id + + +ADMIN_ROUTES = [ + ("get", "/api/admin/quotas"), + ("put", "/api/admin/quotas/instance"), + ("delete", "/api/admin/quotas/instance"), + ("get", "/api/admin/quotas/users/u1"), + ("put", "/api/admin/quotas/users/u1"), + ("delete", "/api/admin/quotas/users/u1"), + ("get", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"), + ("put", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"), + ("delete", "/api/admin/quotas/teams/00000000-0000-0000-0000-000000000000"), +] + + +class TestGuard: + @pytest.mark.parametrize("method, path", ADMIN_ROUTES) + def test_non_admin_forbidden(self, client, db, method, path): + with _as("u1"): + assert getattr(client, method)(path, json={"token_limit": 1}).status_code == 403 + + @pytest.mark.parametrize("method, path", ADMIN_ROUTES) + def test_unauthenticated(self, client, method, path): + with patch("docsgpt.app.handle_auth", return_value=None): + assert getattr(client, method)(path, json={"token_limit": 1}).status_code == 401 + + def test_team_admin_cannot_set_their_teams_allowance(self, client, db): + team_id = _team(db) + db.execute( + text("INSERT INTO team_members (team_id, user_id, role) VALUES (CAST(:t AS uuid), 'lead', 'team_admin')"), + {"t": team_id}, + ) + with _as("lead"): + resp = client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_unlimited": True}) + assert resp.status_code == 403 + assert QuotaPoliciesRepository(db).get("team", team_id) is None + + +class TestInstancePolicy: + def test_set_read_delete(self, client, db): + with _admin(): + put = client.put("/api/admin/quotas/instance", json={"token_limit": 1000, "note": " default "}) + assert put.status_code == 200 + policy = _body(put)["policy"] + assert (policy["token_limit"], policy["note"], policy["updated_by"]) == (1000, "default", "admin1") + + client.put("/api/admin/quotas/instance", json={"bucket": "agent", "cost_limit_usd": 2.5}) + overview = _body(client.get("/api/admin/quotas")) + assert [(p["bucket"], p["token_limit"], p["cost_limit_usd"]) for p in overview["instance"]] == [ + ("all", 1000, None), + ("agent", None, 2.5), + ] + assert overview["period"] == "month" + + assert _body(client.delete("/api/admin/quotas/instance?bucket=agent"))["deleted"] == 1 + assert _body(client.delete("/api/admin/quotas/instance"))["deleted"] == 1 + assert _body(client.get("/api/admin/quotas"))["instance"] == [] + + def test_writes_are_audited(self, client, db): + with _admin(): + client.put("/api/admin/quotas/instance", json={"token_limit": 5}) + client.delete("/api/admin/quotas/instance") + client.delete("/api/admin/quotas/instance") + events = db.execute( + text("SELECT user_id, event, metadata FROM auth_events WHERE event LIKE 'quota_policy_%'") + ).fetchall() + by_event = {e[1]: e for e in events} + # The second delete removed nothing, so it left no event. + assert sorted((e[0], e[1]) for e in events) == [ + ("admin1", "quota_policy_deleted"), + ("admin1", "quota_policy_set"), + ] + metadata = by_event["quota_policy_set"][2] + assert metadata["token_limit"] == 5 and metadata["by"] == "admin1" + + @pytest.mark.parametrize( + "body", + [ + None, + [], + {}, + {"note": "only a note"}, + {"token_limit": -1}, + {"token_limit": 1.5}, + {"token_limit": True}, + {"token_limit": "10"}, + {"token_limit": 2**63}, + {"cost_limit_usd": -0.01}, + {"cost_limit_usd": "5"}, + {"cost_limit_usd": float("inf")}, + {"cost_limit_usd": 1e12}, + {"token_limit": 1, "token_unlimited": True}, + {"cost_limit_usd": 1, "cost_unlimited": True}, + {"token_unlimited": "yes"}, + {"token_limit": 1, "enabled": "no"}, + {"token_limit": 1, "bucket": "everything"}, + {"token_limit": 1, "note": 7}, + ], + ) + def test_invalid_bodies_rejected(self, client, db, body): + with _admin(): + resp = client.put("/api/admin/quotas/instance", json=body) + assert resp.status_code == 400 + assert QuotaPoliciesRepository(db).list_by_scope("instance") == [] + + def test_unknown_bucket_on_delete(self, client, db): + with _admin(): + assert client.delete("/api/admin/quotas/instance?bucket=nope").status_code == 400 + + def test_zero_is_accepted_as_a_block(self, client, db): + with _admin(): + resp = client.put("/api/admin/quotas/instance", json={"token_limit": 0, "cost_limit_usd": 0}) + assert resp.status_code == 200 + assert _body(resp)["policy"]["token_limit"] == 0 + + +class TestTeamPolicy: + def test_set_and_list_with_team_details(self, client, db): + team_id = _team(db, "q-eng", member="u1") + with _admin(): + assert client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_limit": 500}).status_code == 200 + (row,) = _body(client.get("/api/admin/quotas"))["teams"] + assert (row["team_slug"], row["token_limit"], row["member_count"]) == ("q-eng", 500, 1) + assert _body(client.get(f"/api/admin/quotas/teams/{team_id}"))["policies"][0]["token_limit"] == 500 + + @pytest.mark.parametrize("team_id", ["not-a-uuid", "00000000-0000-0000-0000-000000000000"]) + def test_unknown_team(self, client, db, team_id): + with _admin(): + assert client.put(f"/api/admin/quotas/teams/{team_id}", json={"token_limit": 1}).status_code == 404 + assert client.get(f"/api/admin/quotas/teams/{team_id}").status_code == 404 + + +class TestUserPolicy: + def test_unknown_user(self, client, db): + with _admin(): + assert client.put("/api/admin/quotas/users/ghost", json={"token_limit": 1}).status_code == 404 + assert client.get("/api/admin/quotas/users/ghost").status_code == 404 + + def test_effective_limits_name_their_source(self, client, db): + UsersRepository(db).upsert("u1") + small, big = _team(db, "q-small", member="u1"), _team(db, "q-big", member="u1") + repo = QuotaPoliciesRepository(db) + repo.upsert(scope="instance", subject_id=None, token_limit=10, cost_limit_usd=1) + repo.upsert(scope="team", subject_id=small, token_limit=100) + repo.upsert(scope="team", subject_id=big, token_limit=900) + TokenUsageRepository(db).insert(user_id="u1", prompt_tokens=40, cost=0.25) + + with _admin(): + body = _body(client.get("/api/admin/quotas/users/u1")) + overall = body["effective"][0] + assert overall["bucket"] == "all" + assert overall["tokens"] == {"limit": 900.0, "used": 40, "source": "team", "source_id": big} + assert overall["cost"] == {"limit": 1.0, "used": 0.25, "source": "instance", "source_id": None} + assert body["policies"] == [] + + with _admin(): + client.put("/api/admin/quotas/users/u1", json={"token_limit": 50}) + body = _body(client.get("/api/admin/quotas/users/u1")) + assert body["effective"][0]["tokens"]["source"] == "user" + assert body["policies"][0]["token_limit"] == 50 + + def test_user_policy_audit_is_filed_under_the_user(self, client, db): + UsersRepository(db).upsert("u1") + with _admin(): + client.put("/api/admin/quotas/users/u1", json={"cost_unlimited": True}) + row = db.execute( + text("SELECT user_id, metadata FROM auth_events WHERE event = 'quota_policy_set'") + ).one() + assert row[0] == "u1" and row[1]["by"] == "admin1" and row[1]["scope"] == "user" + + +class TestUnpricedModels: + def test_lists_used_catalog_models_without_a_price(self, client, db): + usage = TokenUsageRepository(db) + usage.insert(user_id="u1", prompt_tokens=10, model_id="local-llama") + usage.insert(user_id="u1", prompt_tokens=5, model_id="claude-haiku-4-5", cost=0.1) + usage.insert(user_id="u1", prompt_tokens=7, model_id="7d0c1a52-2f5e-4c53-9a0e-111111111111") + with _admin(), patch("docsgpt.api.admin.quotas.is_priced", lambda m: m == "claude-haiku-4-5"): + unpriced = _body(client.get("/api/admin/quotas"))["unpriced_models"] + assert unpriced == [{"model_id": "local-llama", "tokens": 10, "cost": 0.0}] + + +class TestMyQuota: + def test_unlimited_user_sees_no_buckets(self, client, db): + with _as("u1"): + body = _body(client.get("/api/user/quota")) + assert body == {"success": True, "period": "month", "buckets": []} + + def test_limited_user_sees_usage_without_policy_internals(self, client, db): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100, note="secret note") + TokenUsageRepository(db).insert(user_id="u1", prompt_tokens=30) + with _as("u1"): + body = _body(client.get("/api/user/quota")) + (bucket,) = body["buckets"] + assert bucket["bucket"] == "all" + assert bucket["tokens"] == {"limit": 100.0, "used": 30} + assert bucket["cost"] == {"limit": None, "used": 0.0} + assert "source" not in json.dumps(body) and "secret" not in json.dumps(body) + + def test_a_user_only_sees_their_own_quota(self, client, db): + QuotaPoliciesRepository(db).upsert(scope="user", subject_id="u1", token_limit=100) + with _as("u2"): + assert _body(client.get("/api/user/quota"))["buckets"] == [] + + def test_unauthenticated(self, client): + with patch("docsgpt.app.handle_auth", return_value=None): + assert client.get("/api/user/quota").status_code == 401