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.
This commit is contained in:
Alex committed 2026-09-21 12:11:11 +01:00
1 parent b77561288d
commit 43fad2a865
9 files changed
+415 -3

No files matched your search

+43
View File
@@ -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,
)
@@ -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,