"""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()