host
a18a530bfd
R12 真机：4 线程 × 10 条记账 0 写错，但 verify_chain 返回 (False, first_bad=11)，链 G、链 H 两次复现同一断点；根因是 record() 先 SELECT 链头、再算摘要、再 INSERT 提交，无锁无 事务，两线程读到同一 prev_hash 各写一条，链分叉。修法： 同一实例用 threading.Lock 串行化；SELECT 前 BEGIN IMMEDIATE 持库级写锁（WAL + busy_timeout 下等待而非立即失败），使另一 实例/进程读到的链头必是已提交最新值。record_warning 委托 record()，同款覆盖。新增并发测试：同款 4×10 复测两轮， verify_chain 全真、40 条一条不丢（旧代码下该测试失败）。
648 lines
26 KiB
Python
648 lines
26 KiB
Python
"""Usage metering and tier quota enforcement (P11).
|
||
|
||
Four metered dimensions: tokens per month, messages per day,
|
||
concurrent tasks, storage bytes. Tokens are read from the existing
|
||
``cost_records`` ledger (``cloud/cost.py``) -- the inference channel's
|
||
bookkeeping point -- never re-collected on a second path. Messages,
|
||
task starts/ends, storage snapshots and warnings are recorded in the
|
||
hash-chained ``usage_ledger`` table in this module.
|
||
|
||
Tier limits and prices are a PROJECT PROPOSAL, not user-confirmed
|
||
pricing and not Meta official pricing: the user delegated the numbers
|
||
("set them by market", 2026-11-06) and never confirmed them. Every
|
||
payload this module returns carries ``proposal: True``. Real billing
|
||
(Stripe / payments / invoicing) is NOT in this module and must be
|
||
approved separately; nothing here charges money.
|
||
|
||
Admission semantics: gates are admission-only. A denial returns a
|
||
``quota_exceeded`` dict (HTTP 402 at the server layer); there is no
|
||
API here that terminates an in-flight turn. The pre-existing USD
|
||
quota path in ``cloud/server.py`` still cuts in-flight work at turn
|
||
boundaries via worker polling -- that is a known, documented
|
||
discrepancy (see docs/chv-muse-p11-gap.md §5), not something this
|
||
module silently resolves.
|
||
|
||
Token gate counts ``source="actual"`` rows only, matching the
|
||
``cost.py`` P1 decision ("estimated rows are never real spend and
|
||
must not trip the cap"); estimated tokens are reported separately
|
||
and labelled as estimates.
|
||
"""
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
from datetime import datetime, timezone
|
||
|
||
from cloud.cost import CostStore
|
||
from cloud.db import connect as _db_connect
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier proposal (see module docstring: PROPOSAL, not confirmed pricing).
|
||
# ---------------------------------------------------------------------------
|
||
|
||
GB = 1024 * 1024 * 1024
|
||
|
||
TIER_NAMES = ("free", "paid", "heavy")
|
||
|
||
TIER_LIMITS = {
|
||
"free": {
|
||
"label": "Free",
|
||
"price_usd_per_month": 0,
|
||
"messages_per_day": 50,
|
||
"tokens_per_month": 200000,
|
||
"concurrent_tasks": 1,
|
||
"storage_bytes": 1 * GB,
|
||
},
|
||
"paid": {
|
||
"label": "Paid",
|
||
"price_usd_per_month": 19,
|
||
"messages_per_day": 500,
|
||
"tokens_per_month": 2000000,
|
||
"concurrent_tasks": 5,
|
||
"storage_bytes": 10 * GB,
|
||
},
|
||
"heavy": {
|
||
"label": "Heavy",
|
||
"price_usd_per_month": 99,
|
||
# The proposal never gave Heavy a messages/day number;
|
||
# None = no daily message cap (project decision, not proposal).
|
||
"messages_per_day": None,
|
||
"tokens_per_month": 20000000,
|
||
"concurrent_tasks": 20,
|
||
"storage_bytes": 100 * GB,
|
||
},
|
||
}
|
||
|
||
_LEDGER_KINDS = ("message", "task_start", "task_end", "storage",
|
||
"warning", "usage_query")
|
||
|
||
|
||
def warn_ratio() -> float:
|
||
"""Warning threshold as a fraction of a limit (default 0.8).
|
||
|
||
Env ``CLOUD_USAGE_WARN_RATIO`` overrides; unparsable or
|
||
out-of-range values fall back to 0.8.
|
||
"""
|
||
try:
|
||
value = float(os.environ.get("CLOUD_USAGE_WARN_RATIO", "0.8"))
|
||
except (TypeError, ValueError):
|
||
return 0.8
|
||
if value <= 0 or value > 1:
|
||
return 0.8
|
||
return value
|
||
|
||
|
||
def utc_day(now: float | None = None) -> str:
|
||
"""UTC calendar day (YYYY-MM-DD) for ``now`` (default: current)."""
|
||
stamp = now if now is not None else time.time()
|
||
return datetime.fromtimestamp(stamp, tz=timezone.utc).strftime("%Y-%m-%d")
|
||
|
||
|
||
def utc_month(now: float | None = None) -> str:
|
||
"""UTC calendar month (YYYY-MM) for ``now`` (default: current)."""
|
||
stamp = now if now is not None else time.time()
|
||
return datetime.fromtimestamp(stamp, tz=timezone.utc).strftime("%Y-%m")
|
||
|
||
|
||
def disk_usage_bytes(path: str) -> int:
|
||
"""Sum of file sizes under ``path`` (best effort, bytes).
|
||
|
||
Does not follow symlinks; unreadable files/dirs are skipped.
|
||
Missing path counts as 0.
|
||
"""
|
||
total = 0
|
||
if not path or not os.path.exists(path):
|
||
return 0
|
||
if os.path.isfile(path) and not os.path.islink(path):
|
||
try:
|
||
return os.path.getsize(path)
|
||
except OSError:
|
||
return 0
|
||
for root, dirs, files in os.walk(path, followlinks=False):
|
||
for name in files:
|
||
full = os.path.join(root, name)
|
||
try:
|
||
if not os.path.islink(full):
|
||
total += os.path.getsize(full)
|
||
except OSError:
|
||
continue
|
||
return total
|
||
|
||
|
||
def _entry_hash(prev_hash: str | None, fields: dict) -> str:
|
||
"""Hash-chain digest over the previous hash + canonical fields."""
|
||
payload = json.dumps({"prev": prev_hash or "", "fields": fields},
|
||
sort_keys=True, separators=(",", ":"))
|
||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||
|
||
|
||
class TierStore:
|
||
"""Per-user tier assignment (``user_tiers`` table).
|
||
|
||
No row means the user is NOT tier-gated: only the pre-existing
|
||
USD quota applies. Assigning a tier is an explicit admin action;
|
||
nothing here defaults a user to Free.
|
||
"""
|
||
|
||
_DDL = """CREATE TABLE IF NOT EXISTS user_tiers (
|
||
user_id TEXT PRIMARY KEY,
|
||
tier TEXT,
|
||
updated_at DOUBLE PRECISION
|
||
)"""
|
||
|
||
def __init__(self, db_path: str) -> None:
|
||
self.db_path = db_path
|
||
self._conn = _db_connect(db_path)
|
||
self._conn.execute(self._DDL)
|
||
self._conn.commit()
|
||
|
||
def get_tier(self, user_id: str) -> str | None:
|
||
"""Tier name for ``user_id``, or None when unassigned."""
|
||
row = self._conn.execute(
|
||
"SELECT tier FROM user_tiers WHERE user_id = ?",
|
||
(user_id,)).fetchone()
|
||
if row is None:
|
||
return None
|
||
return row[0]
|
||
|
||
def set_tier(self, user_id: str, tier: str) -> None:
|
||
"""Assign ``tier`` to ``user_id``; unknown tier names rejected."""
|
||
if tier not in TIER_NAMES:
|
||
raise ValueError("unknown tier %r (want one of %r)"
|
||
% (tier, TIER_NAMES))
|
||
self._conn.execute(
|
||
"INSERT INTO user_tiers (user_id, tier, updated_at)"
|
||
" VALUES (?, ?, ?) ON CONFLICT(user_id) DO UPDATE"
|
||
" SET tier = excluded.tier,"
|
||
" updated_at = excluded.updated_at",
|
||
(user_id, tier, time.time()))
|
||
self._conn.commit()
|
||
|
||
def clear_tier(self, user_id: str) -> None:
|
||
"""Remove the tier row: the user is tier-gated no more."""
|
||
self._conn.execute(
|
||
"DELETE FROM user_tiers WHERE user_id = ?", (user_id,))
|
||
self._conn.commit()
|
||
|
||
def limits_for(self, user_id: str) -> dict | None:
|
||
"""Tier limits dict for ``user_id``, or None when unassigned."""
|
||
tier = self.get_tier(user_id)
|
||
if tier is None:
|
||
return None
|
||
return dict(TIER_LIMITS[tier])
|
||
|
||
|
||
class UsageLedger:
|
||
"""Append-only, hash-chained usage ledger (``usage_ledger``).
|
||
|
||
One chain per user: each row stores the previous row's hash and
|
||
its own digest over the previous hash + canonical fields. Any
|
||
UPDATE/DELETE breaks ``verify_chain``. The ledger never stores
|
||
conversation text, file contents, or credential values -- only
|
||
counts, ids, and hashes.
|
||
"""
|
||
|
||
_DDL = """CREATE TABLE IF NOT EXISTS usage_ledger (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
user_id TEXT,
|
||
kind TEXT,
|
||
amount INTEGER,
|
||
unit TEXT,
|
||
task_id TEXT,
|
||
model TEXT,
|
||
source TEXT,
|
||
day TEXT,
|
||
created_at REAL,
|
||
prev_hash TEXT,
|
||
entry_hash TEXT
|
||
)"""
|
||
|
||
def __init__(self, db_path: str) -> None:
|
||
self.db_path = db_path
|
||
if db_path != ":memory:":
|
||
parent = os.path.dirname(os.path.abspath(db_path))
|
||
os.makedirs(parent, exist_ok=True)
|
||
self._conn = _db_connect(db_path)
|
||
self._conn.execute(self._DDL)
|
||
self._conn.commit()
|
||
# round-12 修复轮（项 5，P11）：record() 读-改-写的
|
||
# 串行锁。无锁时两线程读到同一 prev_hash、各写一条，
|
||
# 链分叉（R12 真机：4 线程 × 10 条，0 写错，
|
||
# verify_chain (False, first_bad=11)，链 G、链 H 两次
|
||
# 同断点）。锁只串行化同一实例；跨实例/进程的并发由
|
||
# record() 里的 BEGIN IMMEDIATE 持写锁兜底。
|
||
self._lock = threading.Lock()
|
||
|
||
def _new_conn(self) -> sqlite3.Connection:
|
||
"""Fresh connection per write/read (mirrors CostStore).
|
||
|
||
``:memory:`` reuses the shared connection; anything else
|
||
opens a new one so concurrent threads never share a cursor.
|
||
"""
|
||
if self.db_path == ":memory:":
|
||
return self._conn
|
||
conn = _db_connect(self.db_path)
|
||
conn.execute(self._DDL)
|
||
conn.commit()
|
||
return conn
|
||
|
||
def record(self, user_id: str, kind: str, amount: int,
|
||
unit: str = "", task_id: str | None = None,
|
||
model: str | None = None, source: str | None = None,
|
||
day: str | None = None) -> int:
|
||
"""Append one ledger event; returns its row id.
|
||
|
||
``kind`` must be one of the known event kinds; ``amount`` is
|
||
a count (1 per message, bytes for storage snapshots).
|
||
"""
|
||
if kind not in _LEDGER_KINDS:
|
||
raise ValueError("unknown ledger kind %r" % (kind,))
|
||
day = day or utc_day()
|
||
created_at = time.time()
|
||
conn = self._new_conn()
|
||
try:
|
||
# 读-改-写原子化：锁串行化同一实例；BEGIN IMMEDIATE
|
||
# 在 SELECT 前即持库级写锁（WAL + busy_timeout 下等
|
||
# 待而非立即失败），使另一实例/进程读到的链头必是已
|
||
# 提交的最新值——链不断、条不丢。锁与事务缺一不可：
|
||
# 光有事务、同一实例的两线程仍可能各持一个连接同时读
|
||
# 到同一链头。
|
||
with self._lock:
|
||
conn.execute("BEGIN IMMEDIATE")
|
||
prev = conn.execute(
|
||
"SELECT entry_hash FROM usage_ledger"
|
||
" WHERE user_id = ? ORDER BY id DESC LIMIT 1",
|
||
(user_id,)).fetchone()
|
||
prev_hash = prev[0] if prev is not None else None
|
||
fields = {"user_id": user_id, "kind": kind,
|
||
"amount": int(amount), "unit": unit,
|
||
"task_id": task_id or "", "model": model or "",
|
||
"source": source or "", "day": day,
|
||
"created_at": created_at}
|
||
digest = _entry_hash(prev_hash, fields)
|
||
cursor = conn.execute(
|
||
"""INSERT INTO usage_ledger
|
||
(user_id, kind, amount, unit, task_id, model,
|
||
source, day, created_at, prev_hash, entry_hash)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||
(user_id, kind, int(amount), unit, task_id, model,
|
||
source, day, created_at, prev_hash, digest))
|
||
conn.commit()
|
||
return int(cursor.lastrowid or 0)
|
||
finally:
|
||
if conn is not self._conn:
|
||
conn.close()
|
||
|
||
def sum_kind(self, user_id: str, kind: str,
|
||
day: str | None = None,
|
||
month: str | None = None) -> int:
|
||
"""Sum ``amount`` for ``user_id``/``kind`` in a day or month."""
|
||
conn = self._new_conn()
|
||
try:
|
||
if day is not None:
|
||
row = conn.execute(
|
||
"SELECT COALESCE(SUM(amount), 0) FROM usage_ledger"
|
||
" WHERE user_id = ? AND kind = ? AND day = ?",
|
||
(user_id, kind, day)).fetchone()
|
||
elif month is not None:
|
||
row = conn.execute(
|
||
"SELECT COALESCE(SUM(amount), 0) FROM usage_ledger"
|
||
" WHERE user_id = ? AND kind = ? AND day LIKE ?",
|
||
(user_id, kind, month + "%")).fetchone()
|
||
else:
|
||
row = conn.execute(
|
||
"SELECT COALESCE(SUM(amount), 0) FROM usage_ledger"
|
||
" WHERE user_id = ? AND kind = ?",
|
||
(user_id, kind)).fetchone()
|
||
return int(row[0] or 0)
|
||
finally:
|
||
if conn is not self._conn:
|
||
conn.close()
|
||
|
||
def verify_chain(self, user_id: str) -> tuple[bool, int | None]:
|
||
"""Recompute the hash chain; (True, None) or (False, bad id)."""
|
||
conn = self._new_conn()
|
||
try:
|
||
rows = conn.execute(
|
||
"SELECT id, user_id, kind, amount, unit, task_id,"
|
||
" model, source, day, created_at, prev_hash,"
|
||
" entry_hash FROM usage_ledger"
|
||
" WHERE user_id = ? ORDER BY id",
|
||
(user_id,)).fetchall()
|
||
finally:
|
||
if conn is not self._conn:
|
||
conn.close()
|
||
prev_hash: str | None = None
|
||
for row in rows:
|
||
fields = {"user_id": row[1], "kind": row[2],
|
||
"amount": int(row[3]), "unit": row[4],
|
||
"task_id": row[5] or "", "model": row[6] or "",
|
||
"source": row[7] or "", "day": row[8],
|
||
"created_at": row[9]}
|
||
if row[10] != prev_hash:
|
||
return False, int(row[0])
|
||
if _entry_hash(prev_hash, fields) != row[11]:
|
||
return False, int(row[0])
|
||
prev_hash = row[11]
|
||
return True, None
|
||
|
||
def warning_recorded(self, user_id: str, dimension: str,
|
||
window: str) -> bool:
|
||
"""True when a warning for (dimension, window) is on record."""
|
||
conn = self._new_conn()
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT COUNT(*) FROM usage_ledger"
|
||
" WHERE user_id = ? AND kind = 'warning'"
|
||
" AND unit = ? AND day = ?",
|
||
(user_id, dimension, window)).fetchone()
|
||
return int(row[0] or 0) > 0
|
||
finally:
|
||
if conn is not self._conn:
|
||
conn.close()
|
||
|
||
def record_warning(self, user_id: str, dimension: str,
|
||
window: str) -> int:
|
||
"""Record a once-per-window warning event; returns row id."""
|
||
return self.record(user_id, "warning", 1, unit=dimension,
|
||
day=window)
|
||
|
||
|
||
class UsageMeter:
|
||
"""Four-dimension usage metering for one database.
|
||
|
||
Tokens come from ``cost_records`` (the inference channel's
|
||
bookkeeping point, joined to ``tasks`` exactly like
|
||
``_quota_status`` does); messages come from the ledger;
|
||
concurrency and storage come from injected measurers so the
|
||
meter stays unit-testable without a server or a disk layout.
|
||
An unwired measurer reports ``None`` (unmeasured) in summaries
|
||
and raises in gates -- never silently passes as zero.
|
||
"""
|
||
|
||
def __init__(self, db_path: str,
|
||
running_task_count=None,
|
||
storage_bytes=None) -> None:
|
||
self.db_path = db_path
|
||
self.tiers = TierStore(db_path)
|
||
self.ledger = UsageLedger(db_path)
|
||
self._running_task_count = running_task_count
|
||
self._storage_bytes = storage_bytes
|
||
|
||
def monthly_tokens(self, user_id: str,
|
||
month: str | None = None) -> dict:
|
||
"""Actual / estimated / total tokens for the UTC month.
|
||
|
||
``actual`` is provider-reported usage; ``estimated`` is the
|
||
character-count fallback, labelled as an estimate. Only
|
||
``actual`` ever trips a gate.
|
||
"""
|
||
month = month or utc_month()
|
||
_cost_store = CostStore(self.db_path)
|
||
try:
|
||
_cost_store._conn.close()
|
||
except Exception:
|
||
pass
|
||
conn = _db_connect(self.db_path)
|
||
try:
|
||
rows = conn.execute(
|
||
"SELECT c.source,"
|
||
" COALESCE(SUM(c.input_tokens + c.output_tokens), 0)"
|
||
" FROM cost_records c JOIN tasks t ON c.task_id = t.id"
|
||
" WHERE t.created_by = ? AND c.day LIKE ?"
|
||
" GROUP BY c.source",
|
||
(user_id, month + "%")).fetchall()
|
||
finally:
|
||
conn.close()
|
||
totals = {"actual": 0, "estimated": 0}
|
||
for source, total in rows:
|
||
if source in totals:
|
||
totals[source] = int(total or 0)
|
||
totals["total"] = totals["actual"] + totals["estimated"]
|
||
return totals
|
||
|
||
def daily_messages(self, user_id: str,
|
||
day: str | None = None) -> int:
|
||
"""Messages recorded for the UTC day (fixed window)."""
|
||
return self.ledger.sum_kind(user_id, "message",
|
||
day=day or utc_day())
|
||
|
||
def record_message(self, user_id: str,
|
||
task_id: str | None = None) -> int:
|
||
"""Record one message event; returns the ledger row id."""
|
||
return self.ledger.record(user_id, "message", 1,
|
||
unit="message", task_id=task_id)
|
||
|
||
def record_task_start(self, user_id: str,
|
||
task_id: str | None = None) -> int:
|
||
"""Record a task-start audit event; returns row id."""
|
||
return self.ledger.record(user_id, "task_start", 1,
|
||
unit="task", task_id=task_id)
|
||
|
||
def record_task_end(self, user_id: str,
|
||
task_id: str | None = None) -> int:
|
||
"""Record a task-end audit event; returns row id."""
|
||
return self.ledger.record(user_id, "task_end", 1,
|
||
unit="task", task_id=task_id)
|
||
|
||
def record_storage_snapshot(self, user_id: str,
|
||
size_bytes: int) -> int:
|
||
"""Record a measured storage snapshot; returns row id."""
|
||
return self.ledger.record(user_id, "storage", int(size_bytes),
|
||
unit="bytes")
|
||
|
||
def running_tasks(self, user_id: str) -> int | None:
|
||
"""Concurrent running tasks, or None when unwired."""
|
||
if self._running_task_count is None:
|
||
return None
|
||
return int(self._running_task_count(user_id))
|
||
|
||
def storage_used(self, user_id: str) -> int | None:
|
||
"""Storage bytes, or None when unwired."""
|
||
if self._storage_bytes is None:
|
||
return None
|
||
return int(self._storage_bytes(user_id))
|
||
|
||
def summary(self, user_id: str) -> dict:
|
||
"""Current-period usage, limits, remaining, warnings.
|
||
|
||
Warnings are recorded once per dimension per window (month
|
||
for tokens, day for the rest) and included in the payload.
|
||
"""
|
||
month = utc_month()
|
||
day = utc_day()
|
||
limits = self.tiers.limits_for(user_id)
|
||
tier = self.tiers.get_tier(user_id)
|
||
tokens = self.monthly_tokens(user_id, month)
|
||
messages = self.daily_messages(user_id, day)
|
||
running = self.running_tasks(user_id)
|
||
storage = self.storage_used(user_id)
|
||
ledger_ok, ledger_bad = self.ledger.verify_chain(user_id)
|
||
used = {"tokens_actual": tokens["actual"],
|
||
"tokens_estimated": tokens["estimated"],
|
||
"messages": messages,
|
||
"concurrent_tasks": running,
|
||
"storage_bytes": storage}
|
||
payload: dict = {
|
||
"proposal": True,
|
||
"user_id": user_id,
|
||
"tier": tier,
|
||
"month": month,
|
||
"day": day,
|
||
"used": used,
|
||
"limits": limits,
|
||
"remaining": _remaining(limits, used),
|
||
"ledger_verified": ledger_ok,
|
||
"ledger_first_bad_id": ledger_bad,
|
||
"warnings": [],
|
||
}
|
||
if limits is not None:
|
||
payload["warnings"] = self._check_warnings(
|
||
user_id, limits, used, month, day)
|
||
return payload
|
||
|
||
def _check_warnings(self, user_id: str, limits: dict,
|
||
used: dict, month: str, day: str) -> list:
|
||
"""80%-of-limit warnings, once per dimension per window."""
|
||
ratio = warn_ratio()
|
||
warnings = []
|
||
checks = (("tokens", limits.get("tokens_per_month"),
|
||
used["tokens_actual"], month),
|
||
("messages", limits.get("messages_per_day"),
|
||
used["messages"], day),
|
||
("concurrent_tasks", limits.get("concurrent_tasks"),
|
||
used["concurrent_tasks"], day),
|
||
("storage", limits.get("storage_bytes"),
|
||
used["storage_bytes"], day))
|
||
for dimension, limit, value, window in checks:
|
||
if limit is None or value is None:
|
||
continue
|
||
if limit > 0 and value >= limit * ratio:
|
||
if not self.ledger.warning_recorded(
|
||
user_id, dimension, window):
|
||
self.ledger.record_warning(user_id, dimension,
|
||
window)
|
||
warnings.append({"dimension": dimension,
|
||
"limit": limit, "used": value,
|
||
"ratio": ratio})
|
||
return warnings
|
||
|
||
|
||
def _remaining(limits: dict | None, used: dict) -> dict:
|
||
"""Remaining quota per dimension (None = unlimited/unmeasured)."""
|
||
if limits is None:
|
||
return {"tokens": None, "messages": None,
|
||
"concurrent_tasks": None, "storage_bytes": None}
|
||
pairs = (("tokens", limits.get("tokens_per_month"),
|
||
used["tokens_actual"]),
|
||
("messages", limits.get("messages_per_day"),
|
||
used["messages"]),
|
||
("concurrent_tasks", limits.get("concurrent_tasks"),
|
||
used["concurrent_tasks"]),
|
||
("storage_bytes", limits.get("storage_bytes"),
|
||
used["storage_bytes"]))
|
||
remaining: dict[str, int | None] = {}
|
||
for name, limit, value in pairs:
|
||
if limit is None or value is None:
|
||
remaining[name] = None
|
||
else:
|
||
remaining[name] = max(0, limit - value)
|
||
return remaining
|
||
|
||
|
||
def _denial(dimension: str, limit: int, used: int,
|
||
detail: str) -> dict:
|
||
"""One quota denial payload (HTTP 402 at the server layer)."""
|
||
return {"error": "quota_exceeded", "dimension": dimension,
|
||
"limit": limit, "used": used, "detail": detail}
|
||
|
||
|
||
class QuotaGate:
|
||
"""Admission-only tier gates (pure decisions, no side effects).
|
||
|
||
Every method returns None (allow) or a denial dict. Nothing here
|
||
touches a running task: an in-flight turn is never cut off by
|
||
these gates; a denial only blocks the NEXT admission.
|
||
"""
|
||
|
||
def __init__(self, meter: UsageMeter) -> None:
|
||
self.meter = meter
|
||
|
||
def check_message_admission(self, user_id: str) -> dict | None:
|
||
"""Gate a new message: messages/day, then tokens/month."""
|
||
limits = self.meter.tiers.limits_for(user_id)
|
||
if limits is None:
|
||
return None
|
||
day_limit = limits.get("messages_per_day")
|
||
if day_limit is not None:
|
||
used = self.meter.daily_messages(user_id)
|
||
if used >= day_limit:
|
||
return _denial(
|
||
"messages", day_limit, used,
|
||
"daily message quota exceeded for tier %r"
|
||
" (used %d of %d); wait until tomorrow (UTC)"
|
||
" or upgrade the tier"
|
||
% (self.meter.tiers.get_tier(user_id), used,
|
||
day_limit))
|
||
return self._check_tokens(user_id, limits)
|
||
|
||
def check_task_admission(
|
||
self, user_id: str, running_tasks: int | None = None,
|
||
storage_bytes: int | None = None) -> dict | None:
|
||
"""Gate a new task: tokens, then concurrency, then storage."""
|
||
limits = self.meter.tiers.limits_for(user_id)
|
||
if limits is None:
|
||
return None
|
||
denied = self._check_tokens(user_id, limits)
|
||
if denied is not None:
|
||
return denied
|
||
running = running_tasks
|
||
if running is None:
|
||
running = self.meter.running_tasks(user_id)
|
||
if running is None:
|
||
raise ValueError("no running-task measurer wired")
|
||
task_limit = limits.get("concurrent_tasks")
|
||
if task_limit is not None and running >= task_limit:
|
||
return _denial(
|
||
"concurrent_tasks", task_limit, running,
|
||
"concurrent task quota exceeded for tier %r"
|
||
" (running %d of %d); wait for a running task"
|
||
" to finish" % (self.meter.tiers.get_tier(user_id),
|
||
running, task_limit))
|
||
storage = storage_bytes
|
||
if storage is None:
|
||
storage = self.meter.storage_used(user_id)
|
||
if storage is None:
|
||
raise ValueError("no storage measurer wired")
|
||
storage_limit = limits.get("storage_bytes")
|
||
if storage_limit is not None and storage >= storage_limit:
|
||
return _denial(
|
||
"storage", storage_limit, storage,
|
||
"storage quota exceeded for tier %r"
|
||
" (used %d of %d bytes); free space by removing"
|
||
" old task worktrees or snapshots, then retry"
|
||
% (self.meter.tiers.get_tier(user_id), storage,
|
||
storage_limit))
|
||
return None
|
||
|
||
def _check_tokens(self, user_id: str,
|
||
limits: dict) -> dict | None:
|
||
"""Monthly actual-token gate (estimated tokens never gate)."""
|
||
token_limit = limits.get("tokens_per_month")
|
||
if token_limit is None:
|
||
return None
|
||
used = self.meter.monthly_tokens(user_id)["actual"]
|
||
if used >= token_limit:
|
||
return _denial(
|
||
"tokens", token_limit, used,
|
||
"monthly token quota exceeded for tier %r"
|
||
" (used %d of %d actual tokens; estimated tokens"
|
||
" are not counted); wait until next month (UTC)"
|
||
" or upgrade the tier"
|
||
% (self.meter.tiers.get_tier(user_id), used,
|
||
token_limit))
|
||
return None
|