merge(chv-muse-round12): 合入 P11 用量计量与配额 (60c371d)
- 冲突仅一处:docs/chv-production-readiness.md §15 撞号—— P11 节顺延为 §19(§15 P2 / §16 P4 / §17 P6 / §18 P7 保持不动),节内无自引 §15,无内容改动。 - P11 未决冲突原样保留(见 docs/chv-muse-p11-gap.md §5): 档位门只做准入、USD 路径不动,等待用户裁决,不悄悄统一。 - 测试:test_usage_quota 22/22 OK。
This commit is contained in:
commit
77743e734d
5 files changed
+1474
No files matched your search
+183
@@ -62,6 +62,7 @@ from cloud.auth import AuthStore
|
||||
from cloud.coordinator import parse_parts as _coord_parse_parts
|
||||
from cloud.coordinator import validate_parts as _coord_validate_parts
|
||||
from cloud.cost import CostStore, unpriced_token_cap
|
||||
from cloud.usage import QuotaGate, UsageMeter, disk_usage_bytes
|
||||
from cloud.db import connect as _db_connect
|
||||
from cloud.db import create_database as _create_database
|
||||
from cloud.db import init_postgres_schema as _init_postgres_schema
|
||||
@@ -515,6 +516,52 @@ def _task_worktree_for_files(server, task):
|
||||
return None, (404, "no local worktree for this task")
|
||||
|
||||
|
||||
def _measure_user_storage_bytes(server, user_id):
|
||||
"""Per-user storage measurement (P11): bytes, measured, total.
|
||||
|
||||
Sums each of the user's tasks' local worktrees (via
|
||||
``_task_worktree_for_files``) plus the user's task attachment
|
||||
sizes from the DB. Cloud / cloud-hypervisor tasks have no local
|
||||
worktree and are not measurable on this host: they count in
|
||||
``total`` but not in ``measured``, so callers can say so
|
||||
honestly instead of pretending the number is complete.
|
||||
"""
|
||||
conn = _connect(server.db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
tasks = []
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT id, project_id, executor, worktree_host"
|
||||
" FROM tasks WHERE created_by = ?",
|
||||
(user_id,)).fetchall()
|
||||
tasks = [dict(row) for row in rows]
|
||||
finally:
|
||||
conn.close()
|
||||
total = len(tasks)
|
||||
measured = 0
|
||||
size = 0
|
||||
for task in tasks:
|
||||
worktree, _err = _task_worktree_for_files(server, task)
|
||||
if worktree is None:
|
||||
continue
|
||||
measured += 1
|
||||
size += disk_usage_bytes(worktree)
|
||||
try:
|
||||
conn = _connect(server.db_path)
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT COALESCE(SUM(a.size), 0) FROM task_attachments a"
|
||||
" JOIN tasks t ON a.task_id = t.id"
|
||||
" WHERE t.created_by = ?",
|
||||
(user_id,)).fetchone()
|
||||
size += int(row[0] or 0)
|
||||
finally:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
return size, measured, total
|
||||
|
||||
|
||||
def _safe_worktree_path(worktree, rel):
|
||||
"""worktree 内相对路径的安全解析:空串返回 worktree 本身;
|
||||
绝对路径、.. 逃逸、符号链接逃逸一律返回 None。realpath 后
|
||||
@@ -2294,6 +2341,8 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
self._send_json(404, {"error": "not found"})
|
||||
elif path == "/users/me/quota" and method == "GET":
|
||||
self._user_quota_self(user_id)
|
||||
elif path == "/users/me/usage" and method == "GET":
|
||||
self._user_usage_self(user_id)
|
||||
elif path == "/users/me/profile" and method == "GET":
|
||||
self._identity_profile_get(user_id)
|
||||
elif path == "/users/me/profile" and method in ("PUT", "POST"):
|
||||
@@ -2336,6 +2385,10 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
self._admin_quota_get(user_id, query)
|
||||
elif path == "/admin/quotas" and method == "POST":
|
||||
self._admin_quota_set(user_id)
|
||||
elif path == "/admin/tiers" and method == "GET":
|
||||
self._admin_tiers_get(user_id, query)
|
||||
elif path == "/admin/tiers" and method == "POST":
|
||||
self._admin_tiers_set(user_id)
|
||||
elif path == "/admin/push/dispatch" and method == "POST":
|
||||
self._push_dispatch_endpoint(user_id)
|
||||
elif path.startswith("/tasks/"):
|
||||
@@ -8326,6 +8379,16 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
"unpriced_token_cap": cap,
|
||||
"unpriced_tokens_used": used})
|
||||
return True
|
||||
try:
|
||||
denied = self._tier_gate(quota_user_id or user_id)
|
||||
except Exception:
|
||||
self._send_json(503, {"error": "quota_check_unavailable",
|
||||
"detail": "tier quota lookup failed; "
|
||||
"no task was created or queued"})
|
||||
return True
|
||||
if denied is not None:
|
||||
self._send_json(402, denied)
|
||||
return True
|
||||
return False
|
||||
|
||||
def _reject_if_over_quota_for_task(self, caller_id, task):
|
||||
@@ -8417,6 +8480,126 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
|
||||
payload["unpriced_token_cap"] = unpriced_token_cap()
|
||||
self._send_json(200, payload)
|
||||
|
||||
def _tier_running_tasks(self, user_id):
|
||||
"""该用户正在跑的任务数(queued / running),档位门用。"""
|
||||
conn = self._db()
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT COUNT(*) FROM tasks WHERE created_by = ?"
|
||||
" AND status IN ('queued', 'running')",
|
||||
(user_id,)).fetchone()
|
||||
return int(row[0] or 0)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _tier_meter(self, user_id):
|
||||
"""接好两个实测器的 UsageMeter;存储数只测一次。
|
||||
|
||||
存储实测(磁盘 walk)比较贵:这里先测一次,把字节数与
|
||||
任务计数都交给 meter 的注入函数与 payload,避免一次
|
||||
请求走两遍磁盘。
|
||||
"""
|
||||
measured = {}
|
||||
|
||||
def _storage(uid):
|
||||
if uid not in measured:
|
||||
size, n_measured, n_total = (
|
||||
_measure_user_storage_bytes(self.server, uid))
|
||||
measured[uid] = (size, n_measured, n_total)
|
||||
return measured[uid][0]
|
||||
|
||||
meter = UsageMeter(
|
||||
self.server.db_path,
|
||||
running_task_count=self._tier_running_tasks,
|
||||
storage_bytes=_storage)
|
||||
return meter, measured
|
||||
|
||||
def _tier_gate(self, user_id):
|
||||
"""档位准入:无档位行放行;有则查 token/并发/存储三门。
|
||||
|
||||
返回 None 放行,否则否决 dict(由调用方回 402)。
|
||||
"""
|
||||
meter, _measured = self._tier_meter(user_id)
|
||||
return QuotaGate(meter).check_task_admission(user_id)
|
||||
|
||||
def _user_usage_self(self, user_id):
|
||||
"""GET /users/me/usage:当期用量、剩余额度、提醒、账本校验。
|
||||
|
||||
USD 配额沿用既有 _quota_status 口径并排返回;四项档位
|
||||
用量走 UsageMeter。未定档用户 limits 为 None,只看到
|
||||
用量,看不到档位剩余额度——定档是管理员显式操作。
|
||||
"""
|
||||
meter, measured = self._tier_meter(user_id)
|
||||
payload = meter.summary(user_id)
|
||||
quota, used, month = self._quota_status(user_id)
|
||||
payload["quota_usd_micros"] = quota
|
||||
payload["used_usd_micros"] = used
|
||||
payload["unpriced_tokens_used"] = self._quota_unpriced_tokens(
|
||||
user_id, month)
|
||||
payload["unpriced_token_cap"] = unpriced_token_cap()
|
||||
_size, n_measured, n_total = measured.get(
|
||||
user_id, (payload["used"]["storage_bytes"], 0, 0))
|
||||
payload["storage_measured_tasks"] = n_measured
|
||||
payload["storage_total_tasks"] = n_total
|
||||
self._send_json(200, payload)
|
||||
|
||||
def _admin_tiers_get(self, user_id, query):
|
||||
if not self._is_global_admin(user_id):
|
||||
self._audit_admin(user_id, "admin_tiers_denied",
|
||||
"get tiers")
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
target = (query.get("user_id") or [None])[0]
|
||||
if not target or not isinstance(target, str):
|
||||
self._send_json(400, {"error": "user_id required"})
|
||||
return
|
||||
if not self._user_exists(target):
|
||||
self._send_json(404, {"error": "user not found"})
|
||||
return
|
||||
meter = UsageMeter(self.server.db_path)
|
||||
self._send_json(200, {
|
||||
"proposal": True,
|
||||
"user_id": target,
|
||||
"tier": meter.tiers.get_tier(target),
|
||||
"limits": meter.tiers.limits_for(target)})
|
||||
|
||||
def _admin_tiers_set(self, user_id):
|
||||
if not self._is_global_admin(user_id):
|
||||
self._audit_admin(user_id, "admin_tiers_denied",
|
||||
"set tiers")
|
||||
self._send_json(403, {"error": "forbidden"})
|
||||
return
|
||||
body, err = self._read_json()
|
||||
if err:
|
||||
self._send_json(400, {"error": err})
|
||||
return
|
||||
target = body.get("user_id")
|
||||
if not target or not isinstance(target, str):
|
||||
self._send_json(400, {"error": "user_id required"})
|
||||
return
|
||||
if not self._user_exists(target):
|
||||
self._send_json(404, {"error": "user not found"})
|
||||
return
|
||||
tier = body.get("tier")
|
||||
meter = UsageMeter(self.server.db_path)
|
||||
if tier is None:
|
||||
meter.tiers.clear_tier(target)
|
||||
elif not isinstance(tier, str):
|
||||
self._send_json(400, {"error": "tier must be a string"})
|
||||
return
|
||||
else:
|
||||
try:
|
||||
meter.tiers.set_tier(target, tier)
|
||||
except ValueError:
|
||||
self._send_json(400, {
|
||||
"error": "unknown tier (want free, paid, heavy)"})
|
||||
return
|
||||
self._send_json(200, {
|
||||
"proposal": True,
|
||||
"user_id": target,
|
||||
"tier": meter.tiers.get_tier(target),
|
||||
"limits": meter.tiers.limits_for(target)})
|
||||
|
||||
def _admin_quota_get(self, user_id, query):
|
||||
if not self._is_global_admin(user_id):
|
||||
self._audit_admin(user_id, "admin_quota_denied",
|
||||
|
||||
@@ -0,0 +1,499 @@
|
||||
"""P11 用量计量与配额执行:档位门、计量口径、账本、服务端点。
|
||||
|
||||
不跑真实模型、不联网:本地临时 SQLite + 本地临时端口。
|
||||
档位数值钉住提案(Free $0 / Paid $19 / Heavy $99,未确认定价),
|
||||
token 门只计 actual、estimated 只展示不拦人,账本哈希链防篡改,
|
||||
准入判定不碰在跑任务。
|
||||
"""
|
||||
|
||||
import http.client
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.cost import CostStore
|
||||
from cloud.server import create_server
|
||||
from cloud.usage import (GB, TIER_LIMITS, QuotaGate, TierStore,
|
||||
UsageLedger, UsageMeter, disk_usage_bytes,
|
||||
utc_day)
|
||||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.auth import AuthStore
|
||||
from cloud.cost import CostStore
|
||||
from cloud.server import create_server
|
||||
from cloud.usage import (GB, TIER_LIMITS, QuotaGate, TierStore,
|
||||
UsageLedger, UsageMeter, disk_usage_bytes,
|
||||
utc_day)
|
||||
|
||||
|
||||
def _make_db(tmp):
|
||||
"""临时库文件 + 最小 tasks 表(monthly_tokens 要 JOIN 它)。"""
|
||||
db_path = os.path.join(tmp, "usage.db")
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("CREATE TABLE IF NOT EXISTS tasks ("
|
||||
"id TEXT PRIMARY KEY, created_by TEXT, status TEXT)")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
return db_path
|
||||
|
||||
|
||||
def _seed_tokens(db_path, user_id, actual=0, estimated=0):
|
||||
"""给 user 记一笔 actual 与一笔 estimated(同一任务行)。"""
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("INSERT OR IGNORE INTO tasks (id, created_by, status)"
|
||||
" VALUES ('task-1', ?, 'completed')", (user_id,))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
store = CostStore(db_path)
|
||||
if actual:
|
||||
store.record("task-1", "proj-1", "gpt-6.1-terra",
|
||||
input_tokens=actual, output_tokens=0,
|
||||
source="actual")
|
||||
if estimated:
|
||||
store.record("task-1", "proj-1", "gpt-6.1-terra",
|
||||
input_tokens=estimated, output_tokens=0,
|
||||
source="estimated")
|
||||
|
||||
|
||||
class TierProposalTestCase(unittest.TestCase):
|
||||
def test_proposal_numbers_pinned(self):
|
||||
self.assertEqual(TIER_LIMITS["free"]["price_usd_per_month"], 0)
|
||||
self.assertEqual(TIER_LIMITS["free"]["messages_per_day"], 50)
|
||||
self.assertEqual(TIER_LIMITS["free"]["tokens_per_month"], 200000)
|
||||
self.assertEqual(TIER_LIMITS["free"]["concurrent_tasks"], 1)
|
||||
self.assertEqual(TIER_LIMITS["free"]["storage_bytes"], 1 * GB)
|
||||
self.assertEqual(TIER_LIMITS["paid"]["price_usd_per_month"], 19)
|
||||
self.assertEqual(TIER_LIMITS["paid"]["messages_per_day"], 500)
|
||||
self.assertEqual(TIER_LIMITS["paid"]["tokens_per_month"], 2000000)
|
||||
self.assertEqual(TIER_LIMITS["paid"]["concurrent_tasks"], 5)
|
||||
self.assertEqual(TIER_LIMITS["paid"]["storage_bytes"], 10 * GB)
|
||||
self.assertEqual(TIER_LIMITS["heavy"]["price_usd_per_month"], 99)
|
||||
self.assertIsNone(TIER_LIMITS["heavy"]["messages_per_day"])
|
||||
self.assertEqual(TIER_LIMITS["heavy"]["tokens_per_month"],
|
||||
20000000)
|
||||
self.assertEqual(TIER_LIMITS["heavy"]["concurrent_tasks"], 20)
|
||||
self.assertEqual(TIER_LIMITS["heavy"]["storage_bytes"], 100 * GB)
|
||||
|
||||
def test_unknown_tier_rejected_and_clear(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
store = TierStore(db_path)
|
||||
self.assertIsNone(store.get_tier("u1"))
|
||||
self.assertIsNone(store.limits_for("u1"))
|
||||
with self.assertRaises(ValueError):
|
||||
store.set_tier("u1", "platinum")
|
||||
store.set_tier("u1", "paid")
|
||||
self.assertEqual(store.get_tier("u1"), "paid")
|
||||
self.assertEqual(store.limits_for("u1")["tokens_per_month"],
|
||||
2000000)
|
||||
store.clear_tier("u1")
|
||||
self.assertIsNone(store.get_tier("u1"))
|
||||
|
||||
|
||||
class TokenMeteringTestCase(unittest.TestCase):
|
||||
def test_actual_estimated_split_and_gate_counts_actual_only(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
_seed_tokens(db_path, "u1", actual=1000, estimated=500)
|
||||
tokens = meter.monthly_tokens("u1")
|
||||
self.assertEqual(tokens["actual"], 1000)
|
||||
self.assertEqual(tokens["estimated"], 500)
|
||||
self.assertEqual(tokens["total"], 1500)
|
||||
# estimated 199500 + actual 1000:门只看 actual,不拦
|
||||
_seed_tokens(db_path, "u1", actual=0, estimated=199500)
|
||||
self.assertIsNone(gate.check_message_admission("u1"))
|
||||
# actual 到 200000:拦,维度 tokens
|
||||
_seed_tokens(db_path, "u1", actual=199000, estimated=0)
|
||||
denied = gate.check_message_admission("u1")
|
||||
self.assertIsNotNone(denied)
|
||||
self.assertEqual(denied["error"], "quota_exceeded")
|
||||
self.assertEqual(denied["dimension"], "tokens")
|
||||
self.assertEqual(denied["limit"], 200000)
|
||||
self.assertEqual(denied["used"], 200000)
|
||||
|
||||
def test_unassigned_user_never_tier_gated(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
_seed_tokens(db_path, "u1", actual=99999999)
|
||||
for _ in range(60):
|
||||
meter.record_message("u1")
|
||||
self.assertIsNone(gate.check_message_admission("u1"))
|
||||
self.assertIsNone(gate.check_task_admission(
|
||||
"u1", running_tasks=99, storage_bytes=10 ** 12))
|
||||
|
||||
|
||||
class MessageGateTestCase(unittest.TestCase):
|
||||
def test_daily_fixed_window_and_limit(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
for _ in range(49):
|
||||
meter.record_message("u1")
|
||||
self.assertEqual(meter.daily_messages("u1"), 49)
|
||||
self.assertIsNone(gate.check_message_admission("u1"))
|
||||
meter.record_message("u1")
|
||||
denied = gate.check_message_admission("u1")
|
||||
self.assertEqual(denied["dimension"], "messages")
|
||||
self.assertEqual(denied["limit"], 50)
|
||||
self.assertEqual(denied["used"], 50)
|
||||
# 其它日期的消息不计入今天(用固定过去日期,避免 UTC 日界)
|
||||
meter.ledger.record("u1", "message", 1, day="2026-01-01")
|
||||
self.assertEqual(meter.daily_messages("u1"), 50)
|
||||
self.assertEqual(meter.daily_messages("u1", day="2026-01-01"),
|
||||
1)
|
||||
|
||||
|
||||
class ConcurrencyStorageGateTestCase(unittest.TestCase):
|
||||
def test_concurrency_gate(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
self.assertIsNone(gate.check_task_admission(
|
||||
"u1", running_tasks=0, storage_bytes=0))
|
||||
denied = gate.check_task_admission(
|
||||
"u1", running_tasks=1, storage_bytes=0)
|
||||
self.assertEqual(denied["dimension"], "concurrent_tasks")
|
||||
self.assertEqual(denied["limit"], 1)
|
||||
self.assertIn("wait for a running task", denied["detail"])
|
||||
|
||||
def test_storage_gate_with_cleanup_guidance(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
limit = TIER_LIMITS["free"]["storage_bytes"]
|
||||
self.assertIsNone(gate.check_task_admission(
|
||||
"u1", running_tasks=0, storage_bytes=limit - 1))
|
||||
denied = gate.check_task_admission(
|
||||
"u1", running_tasks=0, storage_bytes=limit)
|
||||
self.assertEqual(denied["dimension"], "storage")
|
||||
self.assertIn("free space", denied["detail"])
|
||||
self.assertIn("worktrees", denied["detail"])
|
||||
|
||||
def test_token_denial_precedes_concurrency(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
_seed_tokens(db_path, "u1", actual=200000)
|
||||
denied = gate.check_task_admission(
|
||||
"u1", running_tasks=99, storage_bytes=10 ** 12)
|
||||
self.assertEqual(denied["dimension"], "tokens")
|
||||
|
||||
def test_unwired_measurer_raises_not_silent_pass(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
with self.assertRaises(ValueError):
|
||||
gate.check_task_admission("u1")
|
||||
|
||||
def test_gate_is_pure_no_ledger_writes(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
gate = QuotaGate(meter)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
meter.record_message("u1")
|
||||
before = meter.ledger.sum_kind("u1", "message")
|
||||
gate.check_message_admission("u1")
|
||||
gate.check_task_admission("u1", running_tasks=5,
|
||||
storage_bytes=10 ** 12)
|
||||
after = meter.ledger.sum_kind("u1", "message")
|
||||
self.assertEqual(before, after)
|
||||
|
||||
|
||||
class DiskUsageTestCase(unittest.TestCase):
|
||||
def test_sums_files_skips_missing_and_links(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self.assertEqual(disk_usage_bytes(
|
||||
os.path.join(tmp, "nope")), 0)
|
||||
sub = os.path.join(tmp, "sub")
|
||||
os.makedirs(sub)
|
||||
with open(os.path.join(tmp, "a.txt"), "wb") as f:
|
||||
f.write(b"x" * 100)
|
||||
with open(os.path.join(sub, "b.txt"), "wb") as f:
|
||||
f.write(b"y" * 50)
|
||||
os.symlink(os.path.join(tmp, "a.txt"),
|
||||
os.path.join(tmp, "link.txt"))
|
||||
self.assertEqual(disk_usage_bytes(tmp), 150)
|
||||
|
||||
|
||||
class LedgerTestCase(unittest.TestCase):
|
||||
def test_chain_verifies_and_tamper_detected(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
ledger = UsageLedger(db_path)
|
||||
ids = [ledger.record("u1", "message", 1)
|
||||
for _ in range(3)]
|
||||
self.assertEqual(ids, sorted(ids))
|
||||
ok, bad = ledger.verify_chain("u1")
|
||||
self.assertTrue(ok)
|
||||
self.assertIsNone(bad)
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("UPDATE usage_ledger SET amount = 999"
|
||||
" WHERE id = ?", (ids[1],))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
ok, bad = ledger.verify_chain("u1")
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(bad, ids[1])
|
||||
|
||||
def test_prev_hash_tamper_detected(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
ledger = UsageLedger(db_path)
|
||||
first = ledger.record("u1", "message", 1)
|
||||
second = ledger.record("u1", "message", 1)
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("UPDATE usage_ledger SET prev_hash = 'x'"
|
||||
" WHERE id = ?", (second,))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
ok, bad = ledger.verify_chain("u1")
|
||||
self.assertFalse(ok)
|
||||
self.assertEqual(bad, second)
|
||||
|
||||
def test_unknown_kind_rejected_and_no_secrets_field(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
ledger = UsageLedger(db_path)
|
||||
with self.assertRaises(ValueError):
|
||||
ledger.record("u1", "chat_text", 1)
|
||||
row_id = ledger.record("u1", "message", 1, task_id="t1")
|
||||
conn = sqlite3.connect(db_path)
|
||||
cols = [r[1] for r in conn.execute(
|
||||
"PRAGMA table_info(usage_ledger)").fetchall()]
|
||||
conn.close()
|
||||
self.assertNotIn("content", cols)
|
||||
self.assertNotIn("token", cols)
|
||||
self.assertGreater(row_id, 0)
|
||||
|
||||
|
||||
class WarningTestCase(unittest.TestCase):
|
||||
def test_once_per_window_at_80_percent(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
for _ in range(39):
|
||||
meter.record_message("u1")
|
||||
self.assertEqual(meter.summary("u1")["warnings"], [])
|
||||
meter.record_message("u1")
|
||||
summary = meter.summary("u1")
|
||||
self.assertEqual(len(summary["warnings"]), 1)
|
||||
self.assertEqual(summary["warnings"][0]["dimension"],
|
||||
"messages")
|
||||
# 同一窗口第二次 summary:提醒仍在,但账本只记一次
|
||||
meter.summary("u1")
|
||||
self.assertEqual(
|
||||
meter.ledger.sum_kind("u1", "warning"), 1)
|
||||
|
||||
def test_warn_ratio_env_override(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
meter.tiers.set_tier("u1", "free")
|
||||
for _ in range(25):
|
||||
meter.record_message("u1")
|
||||
self.assertEqual(meter.summary("u1")["warnings"], [])
|
||||
old = os.environ.get("CLOUD_USAGE_WARN_RATIO")
|
||||
os.environ["CLOUD_USAGE_WARN_RATIO"] = "0.5"
|
||||
try:
|
||||
warnings = meter.summary("u1")["warnings"]
|
||||
finally:
|
||||
if old is None:
|
||||
os.environ.pop("CLOUD_USAGE_WARN_RATIO", None)
|
||||
else:
|
||||
os.environ["CLOUD_USAGE_WARN_RATIO"] = old
|
||||
self.assertEqual(len(warnings), 1)
|
||||
|
||||
|
||||
class SummaryTestCase(unittest.TestCase):
|
||||
def test_summary_shape_and_remaining(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(
|
||||
db_path, running_task_count=lambda uid: 0,
|
||||
storage_bytes=lambda uid: 123)
|
||||
meter.tiers.set_tier("u1", "paid")
|
||||
_seed_tokens(db_path, "u1", actual=1000, estimated=200)
|
||||
meter.record_message("u1")
|
||||
summary = meter.summary("u1")
|
||||
self.assertTrue(summary["proposal"])
|
||||
self.assertEqual(summary["tier"], "paid")
|
||||
self.assertTrue(summary["ledger_verified"])
|
||||
self.assertIsNone(summary["ledger_first_bad_id"])
|
||||
self.assertEqual(summary["used"]["tokens_actual"], 1000)
|
||||
self.assertEqual(summary["used"]["tokens_estimated"], 200)
|
||||
self.assertEqual(summary["used"]["messages"], 1)
|
||||
self.assertEqual(summary["remaining"]["tokens"], 1999000)
|
||||
self.assertEqual(summary["remaining"]["messages"], 499)
|
||||
self.assertEqual(summary["remaining"]["storage_bytes"],
|
||||
10 * GB - 123)
|
||||
|
||||
def test_summary_unassigned_user(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = _make_db(tmp)
|
||||
meter = UsageMeter(db_path)
|
||||
summary = meter.summary("u1")
|
||||
self.assertIsNone(summary["tier"])
|
||||
self.assertIsNone(summary["limits"])
|
||||
self.assertIsNone(summary["remaining"]["tokens"])
|
||||
|
||||
|
||||
class UsageServerTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
self.auth = AuthStore(self.db_path)
|
||||
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
||||
self.outsider_id, self.outsider_key = self.auth.create_user(
|
||||
"outsider")
|
||||
self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
||||
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
||||
self.port = self.server.server_address[1]
|
||||
self.thread = threading.Thread(
|
||||
target=self.server.serve_forever, daemon=True)
|
||||
self.thread.start()
|
||||
status, self.project = self._req(
|
||||
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
|
||||
def tearDown(self):
|
||||
self.server.shutdown()
|
||||
self.server.server_close()
|
||||
self.thread.join(timeout=5)
|
||||
try:
|
||||
self.auth.conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
if self._old_admins is None:
|
||||
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
||||
else:
|
||||
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _req(self, method, path, body=None, key=None):
|
||||
conn = http.client.HTTPConnection(
|
||||
"127.0.0.1", self.port, timeout=10)
|
||||
headers = {}
|
||||
if key:
|
||||
headers["X-API-Key"] = key
|
||||
raw = json.dumps(body) if body is not None else None
|
||||
if raw is not None:
|
||||
headers["Content-Type"] = "application/json"
|
||||
try:
|
||||
conn.request(method, path, body=raw, headers=headers)
|
||||
resp = conn.getresponse()
|
||||
data = resp.read()
|
||||
finally:
|
||||
conn.close()
|
||||
try:
|
||||
parsed = json.loads(data.decode("utf-8")) if data else None
|
||||
except ValueError:
|
||||
parsed = None
|
||||
return resp.status, parsed
|
||||
|
||||
def _set_tier(self, tier, key=None):
|
||||
return self._req("POST", "/admin/tiers",
|
||||
{"user_id": self.owner_id, "tier": tier},
|
||||
key or self.owner_key)
|
||||
|
||||
def test_usage_endpoint_unassigned(self):
|
||||
status, payload = self._req("GET", "/users/me/usage",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(status, 200, payload)
|
||||
self.assertTrue(payload["proposal"])
|
||||
self.assertIsNone(payload["tier"])
|
||||
self.assertIsNone(payload["limits"])
|
||||
self.assertTrue(payload["ledger_verified"])
|
||||
self.assertEqual(payload["used"]["tokens_actual"], 0)
|
||||
|
||||
def test_admin_tiers_set_get_clear_and_auth(self):
|
||||
status, payload = self._set_tier("free")
|
||||
self.assertEqual(status, 200, payload)
|
||||
self.assertEqual(payload["tier"], "free")
|
||||
self.assertEqual(payload["limits"]["messages_per_day"], 50)
|
||||
self.assertTrue(payload["proposal"])
|
||||
status, payload = self._req(
|
||||
"GET", "/admin/tiers?user_id=%s" % self.owner_id,
|
||||
key=self.owner_key)
|
||||
self.assertEqual(status, 200, payload)
|
||||
self.assertEqual(payload["tier"], "free")
|
||||
status, _ = self._set_tier("platinum")
|
||||
self.assertEqual(status, 400)
|
||||
status, _ = self._set_tier("free", key=self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
status, _ = self._req(
|
||||
"GET", "/admin/tiers?user_id=%s" % self.owner_id,
|
||||
key=self.outsider_key)
|
||||
self.assertEqual(status, 403)
|
||||
status, payload = self._set_tier(None)
|
||||
self.assertEqual(status, 200, payload)
|
||||
self.assertIsNone(payload["tier"])
|
||||
|
||||
def test_concurrency_gate_blocks_second_task(self):
|
||||
status, _ = self._set_tier("free")
|
||||
self.assertEqual(status, 200)
|
||||
status, task = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "做事"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201, task)
|
||||
status, denied = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "再做事"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 402, denied)
|
||||
self.assertEqual(denied["error"], "quota_exceeded")
|
||||
self.assertEqual(denied["dimension"], "concurrent_tasks")
|
||||
# 清档后放行
|
||||
status, _ = self._set_tier(None)
|
||||
self.assertEqual(status, 200)
|
||||
status, task2 = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "再做事"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201, task2)
|
||||
|
||||
def test_usage_endpoint_tiered_reflects_running_task(self):
|
||||
self._set_tier("free")
|
||||
status, _ = self._req(
|
||||
"POST", "/tasks",
|
||||
{"project_id": self.project["id"], "prompt": "做事"},
|
||||
self.owner_key)
|
||||
self.assertEqual(status, 201)
|
||||
status, payload = self._req("GET", "/users/me/usage",
|
||||
key=self.owner_key)
|
||||
self.assertEqual(status, 200, payload)
|
||||
self.assertEqual(payload["used"]["concurrent_tasks"], 1)
|
||||
self.assertEqual(payload["remaining"]["concurrent_tasks"], 0)
|
||||
self.assertEqual(payload["storage_total_tasks"], 1)
|
||||
self.assertEqual(payload["storage_measured_tasks"], 0)
|
||||
dims = [w["dimension"] for w in payload["warnings"]]
|
||||
self.assertIn("concurrent_tasks", dims)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+631
@@ -0,0 +1,631 @@
|
||||
"""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
|
||||
Reference in new issue
Block a user