host
60c371d9d5
- cloud/usage.py（新）：TIER_LIMITS 三档提案（Free $0/Paid $19/Heavy $99， 提案非定价，payload 带 proposal 标注）、TierStore（无档位行不受档位门 约束）、UsageLedger（追加写哈希链，篡改可发现）、UsageMeter（token 读 cost_records、消息读账本、并发/存储注入实测）、QuotaGate（纯准入判定， 402，不碰在跑任务）、80% 提醒每维度每窗口一次 - cloud/server.py：_reject_if_over_quota 内加档位门（同一漏斗，双人门 自动覆盖；查询异常 fail-closed 503）、GET /users/me/usage、 GET/POST /admin/tiers、存储实测复用 _task_worktree_for_files - docs/chv-muse-p11-gap.md：差距表（先写表再动手）+ 诚实边界（USD 路径 30 秒轮询掐断与“不掐断”语义不一致是已知待决；并发即拒与设计文档 排队不同；cloud 任务存储量不到；Heavy 消息/天提案未给，本轨定不限） - 测试 22/22，全量 1889 OK（skipped=8）；mypy 在 usage.py 上干净， server.py 5 错均为基线旧错（marker_tools/sub2api/_row_get，与本轨无关） - 真机验证排队在 round-11 与 P4～P9 之后；不接真实计费
632 lines
25 KiB
Python
632 lines
25 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 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()
|
|
|
|
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:
|
|
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
|