host
e4b5545a24
- G1：修正 _check_quota docstring（Local popen 30 秒轮询掐断、 Docker/同步路径只有边界检查），README 配额节同步；followup 创建不查配额是已声明设计，保持不动。 - G2：notify_failure 对插值 shlex.quote，模板错误记日志不静默 吞，命令返回值/超时/异常记 host-run.log；dispatch 前置预算门 与 goal 预算中断补 budget-exceeded 通知（没配命令时只写日志）。 - G5/G6：重试与续跑的配额检查按任务创建者查（费用记创建者名下）， 维护者替他人操作不再用错人的额度；续跑 API 侧补查，超限回 402 且不改状态，不再先 200 排队再悄悄 interrupted。 - 无价上报：配额自查/管理端点返回增加 unpriced_actual_tokens （actual 且 usd_micros 为 NULL，与 CostStore 口径一致），只读 上报、不计价、不拦任务；不给 gpt-5.6-terra 编造价格。 测试：host 68、cloud 760（4 跳过）、deploy 19 全过。
329 lines
13 KiB
Python
329 lines
13 KiB
Python
"""租户配额与计费单测：月度 USD 上限 + 准入 402 + 管理端点。
|
|
|
|
不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
|
|
SQLite；用量直接插 cost_records 行（按任务 created_by 归属、
|
|
UTC day），验证 JOIN 归属、自然月窗口与三路准入（建任务 /
|
|
分支 / 重试）。配额单位是微美元（1 美元 = 1_000_000），
|
|
与 cost_records 一致。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from datetime import date, datetime, timedelta, timezone
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud.auth import AuthStore
|
|
from cloud.server import create_server
|
|
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.server import create_server
|
|
|
|
|
|
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 _set_task_status(db_path, task_id, status):
|
|
conn = sqlite3.connect(db_path)
|
|
try:
|
|
conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
|
|
(status, task_id))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _insert_cost(db_path, task_id, project_id, usd_micros, day):
|
|
conn = sqlite3.connect(db_path)
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO cost_records(task_id, project_id, model,"
|
|
" input_tokens, output_tokens, total_tokens, usd_micros,"
|
|
" source, day, created_at) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
|
(task_id, project_id, "gpt-5.6-terra", 100, 50, 150,
|
|
usd_micros, "actual", day, 0))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _utc_today():
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
|
|
|
|
def _prev_month_day():
|
|
today = date.today()
|
|
first = today.replace(day=1)
|
|
return (first - timedelta(days=1)).isoformat()
|
|
|
|
|
|
class ServerQuotasTestCase(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.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": "qproj"}, self.owner_key)
|
|
self.assertEqual(s, 201, proj)
|
|
self.project_id = proj["id"]
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.tmp.cleanup()
|
|
|
|
def _create_task(self, key=None):
|
|
return _req(self.port, "POST", "/tasks",
|
|
{"project_id": self.project_id,
|
|
"prompt": "quota probe"},
|
|
key if key is not None else self.owner_key)
|
|
|
|
def _set_quota(self, target, quota, key=None):
|
|
return _req(self.port, "POST", "/admin/quotas",
|
|
{"user_id": target, "quota_usd_micros": quota},
|
|
key if key is not None else self.owner_key)
|
|
|
|
def test_self_quota_default_unlimited(self):
|
|
s, q = _req(self.port, "GET", "/users/me/quota",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertIsNone(q["quota_usd_micros"])
|
|
self.assertEqual(q["used_usd_micros"], 0)
|
|
self.assertIsNone(q["remaining_usd_micros"])
|
|
self.assertEqual(q["month"],
|
|
datetime.now(timezone.utc).strftime("%Y-%m"))
|
|
|
|
def test_admin_set_and_get(self):
|
|
s, q = self._set_quota(self.owner_id, 5000000)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertEqual(q["quota_usd_micros"], 5000000)
|
|
s, q = _req(self.port, "GET",
|
|
"/admin/quotas?user_id=" + self.owner_id,
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertEqual(q["quota_usd_micros"], 5000000)
|
|
s, q = _req(self.port, "GET", "/users/me/quota",
|
|
key=self.owner_key)
|
|
self.assertEqual(q["quota_usd_micros"], 5000000)
|
|
# null = 恢复不限额
|
|
s, q = self._set_quota(self.owner_id, None)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertIsNone(q["quota_usd_micros"])
|
|
|
|
def test_admin_set_forbidden_and_validation(self):
|
|
# other 没有任何项目角色，不是管理员
|
|
s, q = self._set_quota(self.owner_id, 100, self.other_key)
|
|
self.assertEqual(s, 403, q)
|
|
s, q = _req(self.port, "GET",
|
|
"/admin/quotas?user_id=" + self.owner_id,
|
|
key=self.other_key)
|
|
self.assertEqual(s, 403, q)
|
|
s, q = self._set_quota(self.owner_id, -1)
|
|
self.assertEqual(s, 400, q)
|
|
s, q = self._set_quota("no-such-user", 100)
|
|
self.assertEqual(s, 404, q)
|
|
s, q = _req(self.port, "POST", "/admin/quotas",
|
|
{"quota_usd_micros": 100}, self.owner_key)
|
|
self.assertEqual(s, 400, q)
|
|
s, q = _req(self.port, "GET", "/admin/quotas",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 400, q)
|
|
|
|
def test_usage_counts_current_month_only(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
5000000, _utc_today())
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
100000000, _prev_month_day())
|
|
s, q = _req(self.port, "GET", "/users/me/quota",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertEqual(q["used_usd_micros"], 5000000)
|
|
|
|
def test_create_blocked_at_quota(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
5000000, _utc_today())
|
|
s, _ = self._set_quota(self.owner_id, 5000000)
|
|
self.assertEqual(s, 200)
|
|
s, body = self._create_task()
|
|
self.assertEqual(s, 402, body)
|
|
self.assertEqual(body["error"], "quota_exceeded")
|
|
self.assertEqual(body["quota_usd_micros"], 5000000)
|
|
self.assertEqual(body["used_usd_micros"], 5000000)
|
|
# 额度调高后放行
|
|
s, _ = self._set_quota(self.owner_id, 6000000)
|
|
self.assertEqual(s, 200)
|
|
s, task2 = self._create_task()
|
|
self.assertEqual(s, 201, task2)
|
|
|
|
def test_fork_blocked_at_quota(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
3000000, _utc_today())
|
|
_set_task_status(self.db_path, task["id"], "completed")
|
|
s, _ = self._set_quota(self.owner_id, 3000000)
|
|
self.assertEqual(s, 200)
|
|
s, body = _req(self.port, "POST",
|
|
"/tasks/%s/fork" % task["id"],
|
|
{"message": "branch probe"}, self.owner_key)
|
|
self.assertEqual(s, 402, body)
|
|
s, _ = self._set_quota(self.owner_id, None)
|
|
self.assertEqual(s, 200)
|
|
s, forked = _req(self.port, "POST",
|
|
"/tasks/%s/fork" % task["id"],
|
|
{"message": "branch probe"}, self.owner_key)
|
|
self.assertEqual(s, 201, forked)
|
|
|
|
def test_retry_blocked_at_quota(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
2000000, _utc_today())
|
|
_set_task_status(self.db_path, task["id"], "failed")
|
|
s, _ = self._set_quota(self.owner_id, 2000000)
|
|
self.assertEqual(s, 200)
|
|
s, body = _req(self.port, "POST",
|
|
"/tasks/%s/retry" % task["id"], {},
|
|
self.owner_key)
|
|
self.assertEqual(s, 402, body)
|
|
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
|
key=self.owner_key)
|
|
self.assertEqual(row["status"], "failed")
|
|
s, _ = self._set_quota(self.owner_id, None)
|
|
self.assertEqual(s, 200)
|
|
s, retried = _req(self.port, "POST",
|
|
"/tasks/%s/retry" % task["id"], {},
|
|
self.owner_key)
|
|
self.assertEqual(s, 200, retried)
|
|
self.assertEqual(retried["status"], "queued")
|
|
|
|
def test_retry_uses_creator_quota_not_operator(self):
|
|
# 重试按任务创建者查配额：费用记创建者名下，操作人
|
|
# 是谁不应影响准入。other 须是 maintainer 才能过
|
|
# approve_merge 权限门，否则 403 会掩盖配额判断。
|
|
self.auth.set_member(self.project_id, self.owner_id,
|
|
self.other_id, "maintainer")
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
2000000, _utc_today())
|
|
_set_task_status(self.db_path, task["id"], "failed")
|
|
# 创建者超限：维护者替他重试也应拦。
|
|
s, _ = self._set_quota(self.owner_id, 2000000)
|
|
self.assertEqual(s, 200)
|
|
s, body = _req(self.port, "POST",
|
|
"/tasks/%s/retry" % task["id"], {},
|
|
self.other_key)
|
|
self.assertEqual(s, 402, body)
|
|
# 创建者不限额、操作人额度为 0：按旧口径（查操作人）
|
|
# 会拦，按创建者口径应放行。
|
|
s, _ = self._set_quota(self.owner_id, None)
|
|
self.assertEqual(s, 200)
|
|
s, _ = self._set_quota(self.other_id, 0)
|
|
self.assertEqual(s, 200)
|
|
s, retried = _req(self.port, "POST",
|
|
"/tasks/%s/retry" % task["id"], {},
|
|
self.other_key)
|
|
self.assertEqual(s, 200, retried)
|
|
self.assertEqual(retried["status"], "queued")
|
|
|
|
def test_continue_blocked_at_creator_quota(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
5000000, _utc_today())
|
|
_set_task_status(self.db_path, task["id"], "completed")
|
|
s, _ = self._set_quota(self.owner_id, 5000000)
|
|
self.assertEqual(s, 200)
|
|
# 以前：先回 200 排队、随后认领门悄悄标 interrupted。
|
|
# 现在：API 侧直接 402，状态保持 completed 不变。
|
|
s, body = _req(self.port, "POST",
|
|
"/tasks/%s/continue" % task["id"],
|
|
{"message": "again"}, self.owner_key)
|
|
self.assertEqual(s, 402, body)
|
|
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
|
key=self.owner_key)
|
|
self.assertEqual(row["status"], "completed")
|
|
# 额度调高后放行。
|
|
s, _ = self._set_quota(self.owner_id, None)
|
|
self.assertEqual(s, 200)
|
|
s, continued = _req(self.port, "POST",
|
|
"/tasks/%s/continue" % task["id"],
|
|
{"message": "again"}, self.owner_key)
|
|
self.assertEqual(s, 200, continued)
|
|
self.assertEqual(continued["status"], "queued")
|
|
# 续跑同样按创建者查：操作人额度为 0、创建者不限额时放行。
|
|
self.auth.set_member(self.project_id, self.owner_id,
|
|
self.other_id, "maintainer")
|
|
s, task2 = self._create_task()
|
|
self.assertEqual(s, 201, task2)
|
|
_set_task_status(self.db_path, task2["id"], "completed")
|
|
s, _ = self._set_quota(self.other_id, 0)
|
|
self.assertEqual(s, 200)
|
|
s, continued2 = _req(self.port, "POST",
|
|
"/tasks/%s/continue" % task2["id"],
|
|
{"message": "again"}, self.other_key)
|
|
self.assertEqual(s, 200, continued2)
|
|
self.assertEqual(continued2["status"], "queued")
|
|
|
|
def test_quota_reports_unpriced_tokens(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
# 有价 5000000 计入 used；无价 actual 150 只进
|
|
# unpriced_actual_tokens，不计价、不拦任务。
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
5000000, _utc_today())
|
|
_insert_cost(self.db_path, task["id"], self.project_id,
|
|
None, _utc_today())
|
|
s, q = _req(self.port, "GET", "/users/me/quota",
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertEqual(q["used_usd_micros"], 5000000)
|
|
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
|
s, q = _req(self.port, "GET",
|
|
"/admin/quotas?user_id=" + self.owner_id,
|
|
key=self.owner_key)
|
|
self.assertEqual(s, 200, q)
|
|
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|