host
488aedbcad
(cherry picked from commit 1db02545f74c96b3b92be4fc32954659cec1f814) Reconcile-conflict: A14b E402 imports referencing cloud.handler_core (does not exist here; core methods live in server.py) — kept HEAD imports with incoming's noqa: E402 markers, dropped the unused core_mod import (verified unused elsewhere in both files).
400 lines
17 KiB
Python
400 lines
17 KiB
Python
"""成本与配额成套单测：双列报告、未定价单独列、按创建者分账、
|
|
超限硬切断、审批超时调阈三档。
|
|
|
|
不跑真实模型回合、不联网、不登录任何服务：成本行直接写
|
|
CostStore，配额判定走 quota_cutoff_reason 与 worker 的
|
|
_quota_exceeded（假 runner，不 dispatch），服务端点走本地
|
|
临时端口 + 临时 SQLite。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
try:
|
|
from cloud.approval_timeout import (
|
|
DEFAULT_TIMEOUT_SECONDS, RAISED_TIMEOUT_SECONDS,
|
|
evaluate_approval_timeout, recommend_approval_timeout)
|
|
from cloud.approvals import (
|
|
create_approval, decide_approval, expire_overdue)
|
|
from cloud.auth import AuthStore
|
|
from cloud.cost import (
|
|
BILLING_NOTE, CostStore, price_usd_micros, quota_cutoff_reason)
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.approval_timeout import (
|
|
DEFAULT_TIMEOUT_SECONDS, RAISED_TIMEOUT_SECONDS,
|
|
evaluate_approval_timeout, recommend_approval_timeout)
|
|
from cloud.approvals import (
|
|
create_approval, decide_approval, expire_overdue)
|
|
from cloud.auth import AuthStore
|
|
from cloud.cost import (
|
|
BILLING_NOTE, CostStore, price_usd_micros, quota_cutoff_reason)
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
|
|
|
|
def _req(port, method, path, body=None, key=None):
|
|
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10)
|
|
hdrs = {}
|
|
payload = None
|
|
if body is not None:
|
|
payload = json.dumps(body).encode("utf-8")
|
|
hdrs["Content-Type"] = "application/json"
|
|
if key is not None:
|
|
hdrs["X-API-Key"] = key
|
|
try:
|
|
conn.request(method, path, body=payload, headers=hdrs)
|
|
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 _insert_task_row(db_path, task_id, created_by, project_id="p1"):
|
|
now = time.time()
|
|
conn = sqlite3.connect(db_path)
|
|
try:
|
|
conn.execute(
|
|
"CREATE TABLE IF NOT EXISTS tasks(id TEXT PRIMARY KEY,"
|
|
" project_id TEXT, prompt TEXT, status TEXT,"
|
|
" created_by TEXT, created_at REAL, updated_at REAL)")
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" created_by, created_at, updated_at)"
|
|
" VALUES(?, ?, 'work', 'completed', ?, ?, ?)",
|
|
(task_id, project_id, created_by, now, now))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class DualColumnTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
self.store = CostStore(self.db_path)
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def _seed(self):
|
|
# 有价模型：gpt-6.1-sol 按价格表计价。
|
|
micros = price_usd_micros("gpt-6.1-sol", 10000, 0, 1000)
|
|
self.assertIsNotNone(micros)
|
|
self.store.record("t-priced", "p1", "gpt-6.1-sol",
|
|
input_tokens=10000, output_tokens=1000,
|
|
usd_micros=micros, source="actual")
|
|
# 未定价模型：usd_micros=None，绝不是 0 美元。
|
|
self.store.record("t-free", "p1", "gpt-5.6-terra",
|
|
input_tokens=300, output_tokens=200,
|
|
usd_micros=None, source="actual")
|
|
# 估算行不进未定价 actual 口径。
|
|
self.store.record("t-est", "p1", "gpt-5.6-terra",
|
|
input_tokens=999, output_tokens=0,
|
|
usd_micros=None, source="estimated")
|
|
return micros
|
|
|
|
def test_dual_column_and_unpriced_listed_separately(self):
|
|
micros = self._seed()
|
|
report = self.store.dual_column_report_by_project("p1")
|
|
self.assertEqual(report["priced_usd_micros"], micros)
|
|
self.assertEqual(report["usd_micros"], micros)
|
|
# 未定价 actual 只有 500（估算 999 不计），且单独列出。
|
|
self.assertEqual(report["unpriced_actual_tokens"], 500)
|
|
self.assertEqual(report["unpriced_models"], ["gpt-5.6-terra"])
|
|
self.assertIn("非中转实际账单", report["billing_note"])
|
|
self.assertEqual(report["billing_note"], BILLING_NOTE)
|
|
by_model = {row["model"]: row for row in report["by_model"]}
|
|
self.assertEqual(by_model["gpt-6.1-sol"]["priced_usd_micros"],
|
|
micros)
|
|
self.assertTrue(by_model["gpt-6.1-sol"]["priced"])
|
|
# 未定价行的美元是 None，不是 0：合计绝不把它按免费算。
|
|
self.assertIsNone(by_model["gpt-5.6-terra"]["priced_usd_micros"])
|
|
self.assertFalse(by_model["gpt-5.6-terra"]["priced"])
|
|
self.assertEqual(by_model["gpt-5.6-terra"]["unpriced_tokens"],
|
|
500)
|
|
|
|
def test_task_dual_column_keeps_old_keys(self):
|
|
micros = self._seed()
|
|
report = self.store.dual_column_report_by_task("t-priced")
|
|
self.assertEqual(report["total_tokens"], 11000)
|
|
self.assertEqual(report["usd_micros"], micros)
|
|
self.assertEqual(report["unpriced_actual_tokens"], 0)
|
|
self.assertEqual(report["unpriced_models"], [])
|
|
self.assertEqual(report["turns"], 1)
|
|
self.assertEqual(report["llm_requests"], 0)
|
|
|
|
def test_model_breakdown_requires_exactly_one_scope(self):
|
|
with self.assertRaises(ValueError):
|
|
self.store.model_breakdown()
|
|
with self.assertRaises(ValueError):
|
|
self.store.model_breakdown(task_id="a", project_id="p1")
|
|
with self.assertRaises(ValueError):
|
|
self.store.report_by_creator("u1", month="2026/10")
|
|
|
|
|
|
class CreatorSplitTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
self.store = CostStore(self.db_path)
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_report_by_creator_only_counts_own_tasks(self):
|
|
_insert_task_row(self.db_path, "t-a", "alice")
|
|
_insert_task_row(self.db_path, "t-b", "bob")
|
|
_insert_task_row(self.db_path, "t-orphan", None)
|
|
self.store.record("t-a", "p1", "gpt-6.1-sol",
|
|
input_tokens=100, output_tokens=10,
|
|
usd_micros=1000, source="actual")
|
|
self.store.record("t-b", "p1", "gpt-6.1-sol",
|
|
input_tokens=200, output_tokens=20,
|
|
usd_micros=2000, source="actual")
|
|
self.store.record("t-orphan", "p1", "gpt-6.1-sol",
|
|
input_tokens=999, output_tokens=0,
|
|
usd_micros=9999, source="actual")
|
|
alice = self.store.report_by_creator("alice")
|
|
self.assertEqual(alice["usd_micros"], 1000)
|
|
self.assertEqual(alice["total_tokens"], 110)
|
|
self.assertEqual(alice["creator_id"], "alice")
|
|
bob = self.store.dual_column_report_by_creator("bob")
|
|
self.assertEqual(bob["usd_micros"], 2000)
|
|
# 无属主任务的成本不归任何人，不会漏进 alice/bob。
|
|
self.assertEqual(self.store.report_by_creator("nobody")[
|
|
"usd_micros"], 0)
|
|
month = datetime.now(timezone.utc).strftime("%Y-%m")
|
|
self.assertEqual(
|
|
self.store.report_by_creator("alice", month=month)[
|
|
"usd_micros"], 1000)
|
|
self.assertEqual(
|
|
self.store.report_by_creator("alice", month="2000-01")[
|
|
"usd_micros"], 0)
|
|
|
|
|
|
class HardCutoffTestCase(unittest.TestCase):
|
|
def test_quota_cutoff_reason_three_cases(self):
|
|
# USD 超限。
|
|
self.assertIn("quota_exceeded",
|
|
quota_cutoff_reason(100000, 100000, 0, "2026-10"))
|
|
self.assertIn("quota_exceeded",
|
|
quota_cutoff_reason(100000, 150000, 0, "2026-10"))
|
|
# 无价上限超限（USD 额度还很大也不放行）。
|
|
self.assertIn("quota_exceeded",
|
|
quota_cutoff_reason(10 ** 12, 0, 500, "2026-10",
|
|
unpriced_cap=500))
|
|
# 未超限与显式不限额都不挡。
|
|
self.assertIsNone(quota_cutoff_reason(100000, 99999, 0,
|
|
"2026-10"))
|
|
self.assertIsNone(quota_cutoff_reason(None, 10 ** 9, 10 ** 9,
|
|
"2026-10"))
|
|
|
|
def test_worker_quota_exceeded_counts_in_flight_usage(self):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
db_path = os.path.join(tmp.name, "control.db")
|
|
init_db(db_path)
|
|
conn = sqlite3.connect(db_path)
|
|
try:
|
|
now = time.time()
|
|
conn.execute(
|
|
"INSERT INTO user_quotas(user_id, quota_usd_micros,"
|
|
" updated_at) VALUES(?, ?, ?)",
|
|
("tester", 100000, now))
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" idempotency_key, depends_on, priority, created_by,"
|
|
" created_at, updated_at, last_heartbeat)"
|
|
" VALUES(?, 'p1', 'work', 'queued', NULL, '[]', 0,"
|
|
" 'tester', ?, ?, NULL)",
|
|
("T1", now, now))
|
|
conn.execute(
|
|
"INSERT INTO tasks(id, project_id, prompt, status,"
|
|
" idempotency_key, depends_on, priority, created_by,"
|
|
" created_at, updated_at, last_heartbeat)"
|
|
" VALUES(?, 'p1', 'work', 'completed', NULL, '[]', 0,"
|
|
" 'tester', ?, ?, NULL)",
|
|
("donor", now, now))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
store = CostStore(db_path)
|
|
store.record("donor", "p1", "gpt-6.1-sol",
|
|
input_tokens=1, output_tokens=0,
|
|
usd_micros=90000, source="actual")
|
|
worker = Worker(db_path, "/fake/host.py",
|
|
runner=lambda cmd: SimpleNamespace(
|
|
returncode=0, stdout="{}", stderr=""),
|
|
env={}, model="gpt-6.1-sol",
|
|
retry_backoff_seconds=0)
|
|
task = {"id": "T1", "created_by": "tester",
|
|
"model": "gpt-6.1-sol"}
|
|
# 本轮已花 90000；再花 10000 input + 1000 output =
|
|
# 30000 微美元，合计 120000 >= 100000，轮次边界必须停。
|
|
err = worker._quota_exceeded(
|
|
task, {"input_tokens": 10000, "cached_input_tokens": 0,
|
|
"output_tokens": 1000, "total_tokens": 11000})
|
|
self.assertIsNotNone(err)
|
|
self.assertIn("quota_exceeded", err)
|
|
# 小用量不挡。
|
|
self.assertIsNone(worker._quota_exceeded(
|
|
task, {"input_tokens": 100, "cached_input_tokens": 0,
|
|
"output_tokens": 10, "total_tokens": 110}))
|
|
|
|
|
|
class ApprovalTimeoutTestCase(unittest.TestCase):
|
|
def test_recommend_three_tiers_and_boundaries(self):
|
|
# >20% 调 30 分钟。
|
|
rec = recommend_approval_timeout(3, 10)
|
|
self.assertEqual(rec["action"], "raise_to_30m")
|
|
self.assertEqual(rec["recommended_seconds"],
|
|
RAISED_TIMEOUT_SECONDS)
|
|
self.assertEqual(rec["recommended_seconds"], 1800)
|
|
# <5% 不动。
|
|
rec = recommend_approval_timeout(1, 100)
|
|
self.assertEqual(rec["action"], "keep")
|
|
self.assertEqual(rec["recommended_seconds"],
|
|
DEFAULT_TIMEOUT_SECONDS)
|
|
self.assertEqual(rec["recommended_seconds"], 600)
|
|
# 5%~20% 之间不动；边界值本身也不动。
|
|
self.assertEqual(recommend_approval_timeout(1, 10)["action"],
|
|
"keep")
|
|
self.assertEqual(recommend_approval_timeout(2, 10)["action"],
|
|
"keep")
|
|
self.assertEqual(recommend_approval_timeout(5, 100)["action"],
|
|
"keep")
|
|
# 无已终态时算不出来，不是 0%，同样不动。
|
|
rec = recommend_approval_timeout(0, 0)
|
|
self.assertIsNone(rec["timeout_rate"])
|
|
self.assertEqual(rec["action"], "keep")
|
|
|
|
def test_evaluate_counts_expired_and_logs(self):
|
|
tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(tmp.cleanup)
|
|
conn = sqlite3.connect(os.path.join(tmp.name, "appr.db"))
|
|
conn.row_factory = sqlite3.Row
|
|
try:
|
|
items = []
|
|
for i in range(10):
|
|
items.append(create_approval(
|
|
conn, "t1", "p1", summary="probe %d" % i,
|
|
timeout_seconds=60, now=1000.0))
|
|
for item in items[:7]:
|
|
decide_approval(conn, item["id"], "approved", "tester")
|
|
# 剩余 3 条过期：扫荡后 status=expired。
|
|
self.assertEqual(expire_overdue(conn, now=2000.0), 3)
|
|
with self.assertLogs("cloud.approval_timeout",
|
|
level="INFO") as logs:
|
|
result = evaluate_approval_timeout(conn)
|
|
self.assertEqual(result["decided_count"], 10)
|
|
self.assertEqual(result["expired_count"], 3)
|
|
self.assertAlmostEqual(result["timeout_rate"], 0.3)
|
|
self.assertEqual(result["action"], "raise_to_30m")
|
|
self.assertEqual(result["recommended_seconds"], 1800)
|
|
self.assertTrue(any("approval_timeout metric" in line
|
|
for line in logs.output))
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
class ServerCostEndpointsTestCase(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.other_id, self.other_key = self.auth.create_user("other")
|
|
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()
|
|
s, proj = _req(self.port, "POST", "/projects",
|
|
{"name": "cproj"}, self.owner_key)
|
|
self.assertEqual(s, 201, proj)
|
|
self.project_id = proj["id"]
|
|
s, task = _req(self.port, "POST", "/tasks",
|
|
{"project_id": self.project_id,
|
|
"prompt": "cost probe"}, self.owner_key)
|
|
self.assertEqual(s, 201, task)
|
|
self.task_id = task["id"]
|
|
store = CostStore(self.db_path)
|
|
store.record(self.task_id, self.project_id, "gpt-6.1-sol",
|
|
input_tokens=10000, output_tokens=1000,
|
|
usd_micros=30000, source="actual")
|
|
store.record(self.task_id, self.project_id, "gpt-5.6-terra",
|
|
input_tokens=300, output_tokens=200,
|
|
usd_micros=None, source="actual")
|
|
try:
|
|
store._conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.tmp.cleanup()
|
|
if self._old_admins is None:
|
|
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
|
else:
|
|
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
|
|
|
def test_user_cost_and_task_cost_dual_column(self):
|
|
s, body = _req(self.port, "GET", "/users/me/cost",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, body)
|
|
self.assertEqual(body["cost"]["priced_usd_micros"], 30000)
|
|
self.assertEqual(body["cost"]["unpriced_actual_tokens"], 500)
|
|
self.assertEqual(body["cost"]["unpriced_models"],
|
|
["gpt-5.6-terra"])
|
|
self.assertIn("非中转实际账单",
|
|
body["cost"]["billing_note"])
|
|
s, body = _req(self.port, "GET",
|
|
"/tasks/%s/cost" % self.task_id,
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, body)
|
|
self.assertEqual(body["priced_usd_micros"], 30000)
|
|
self.assertEqual(body["unpriced_models"], ["gpt-5.6-terra"])
|
|
s, body = _req(self.port, "GET",
|
|
"/users/me/cost?month=2026/10",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 400, body)
|
|
|
|
def test_admin_approval_timeout_ops_gated(self):
|
|
s, body = _req(self.port, "GET",
|
|
"/admin/ops/approval-timeout",
|
|
key=self.other_key)
|
|
self.assertEqual(s, 403, body)
|
|
s, body = _req(self.port, "GET",
|
|
"/admin/ops/approval-timeout",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, body)
|
|
self.assertEqual(body["current_seconds"], 600)
|
|
self.assertEqual(body["decided_count"], 0)
|
|
self.assertIsNone(body["timeout_rate"])
|
|
self.assertEqual(body["action"], "keep")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|