host
7e081211db
- POST /admin/sweep 改走任一项目 maintainer+ 门槛（与备份管理同口径），否则 403，记 admin_sweep 审计；threshold 下限钳 60 秒 - OAuth 建号 created=True 时写默认配额行（默认 $10/月，CLOUD_DEFAULT_QUOTA_USD_MICROS 可调，INSERT OR IGNORE 幂等）；老用户无配额行不受影响 - 建任务未指定且无环境继承时默认 executor=docker（分支继承为空时同理）；CLOUD_DEFAULT_EXECUTOR 可覆盖，空字符串恢复旧行为；命令台下拉默认改为 docker - 无价模型 token 上限：价目表外模型（默认 terra）成本 usd 为 NULL，USD 配额永不生效；有配额行的用户当月无价 token 达 CLOUD_QUOTA_UNPRICED_TOKEN_CAP（默认 2000 万/月）即停，服务端 402、worker 认领/重试/后续消息同口径；报表仍记 unpriced，不折假美元 - 测试：新增 test_multiuser_gaps 11 个 + test_worker_quota 3 个；全量 cloud 668、host 64、deploy 19 全过
336 lines
13 KiB
Python
336 lines
13 KiB
Python
"""cloud worker 配额加固单测：认领门、后续消息执行时查配额、轮询掐断。
|
||
|
||
运行（在 ~/workspace/codex-app-server 下）：
|
||
python3 -m unittest discover -s cloud/tests -p "test_worker_quota.py" -v
|
||
|
||
说明：只用注入的假 runner（假 CLI）与真 sleep 子进程，不跑真实
|
||
模型回合、不联网、不登录任何服务；env 映射也是注入的普通
|
||
dict，里面的值只是占位用的假名，绝不是真实密钥。
|
||
|
||
价目（cloud/cost.py，gpt-6.1-sol：input $2.00 / cached $0.10 /
|
||
output $10.00 每百万 token）：input 10000 + output 1000 =
|
||
20000 + 10000 = 30000 微美元。配额 100000 微美元（$0.10）。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
try:
|
||
from cloud.cost import CostStore
|
||
from cloud.server import 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.cost import CostStore
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker
|
||
|
||
QUOTA = 100000
|
||
|
||
|
||
def _insert_task(db_path, task_id, created_by="tester", status="queued",
|
||
project_id="p1"):
|
||
now = time.time()
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" idempotency_key, depends_on, priority, created_by,"
|
||
" created_at, updated_at, last_heartbeat)"
|
||
" VALUES(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)",
|
||
(task_id, project_id, "do work", status, created_by, now, now),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _insert_quota(db_path, user_id, micros):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO user_quotas(user_id, quota_usd_micros, updated_at)"
|
||
" VALUES(?, ?, ?)",
|
||
(user_id, micros, time.time()),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _insert_donor_cost(db_path, micros):
|
||
"""给 tester 记一笔已花成本：另起一个已完成任务并记 actual 行。"""
|
||
_insert_task(db_path, "donor", status="completed")
|
||
store = CostStore(db_path)
|
||
try:
|
||
store.record("donor", "p1", "gpt-6.1-sol", input_tokens=1,
|
||
output_tokens=0, usd_micros=micros, source="actual")
|
||
finally:
|
||
try:
|
||
store._conn.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _insert_unpriced_cost(db_path, tokens):
|
||
"""给 tester 记一笔无价成本：已完成任务 + usd_micros NULL 的行。
|
||
|
||
tokens 拆成 input/output 两半记（input + output = tokens，
|
||
与服务端/ worker 的无价 token 口径一致）。
|
||
"""
|
||
_insert_task(db_path, "donor-unpriced", status="completed")
|
||
store = CostStore(db_path)
|
||
try:
|
||
store.record("donor-unpriced", "p1", "gpt-5.6-terra",
|
||
input_tokens=tokens // 2,
|
||
output_tokens=tokens - tokens // 2,
|
||
usd_micros=None, source="actual")
|
||
finally:
|
||
try:
|
||
store._conn.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _insert_followup(db_path, fu_id, task_id, body="follow up"):
|
||
now = time.time()
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO task_followups(id, task_id, project_id, body,"
|
||
" mode, status, created_by, created_at, updated_at)"
|
||
" VALUES(?, ?, 'p1', ?, 'queued', 'pending', 'tester', ?, ?)",
|
||
(fu_id, task_id, body, now, now),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _status(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT status FROM tasks WHERE id = ?", (task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return row[0] if row else None
|
||
|
||
|
||
def _error_text(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT error_text FROM tasks WHERE id = ?",
|
||
(task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return row[0] if row else None
|
||
|
||
|
||
def _followup_status(db_path, fu_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT status FROM task_followups WHERE id = ?",
|
||
(fu_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return row[0] if row else None
|
||
|
||
|
||
def _cost_rows(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
rows = conn.execute(
|
||
"SELECT usd_micros, source FROM cost_records"
|
||
" WHERE task_id = ?", (task_id,)).fetchall()
|
||
finally:
|
||
conn.close()
|
||
return rows
|
||
|
||
|
||
def _usage_runner(usage, calls):
|
||
"""假 CLI：dispatch/run 一律成功，stdout JSON 里带 usage。"""
|
||
def fake(cmd):
|
||
calls.append(list(cmd))
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout=json.dumps({"status": "completed", "text": "done",
|
||
"usage": usage}),
|
||
stderr="")
|
||
return fake
|
||
|
||
|
||
class WorkerQuotaTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
self.env = {"XCM_API_KEY": "test-placeholder"}
|
||
self.calls = []
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _worker(self, runner):
|
||
return Worker(self.db_path, "/fake/host.py", runner=runner,
|
||
env=self.env, model="gpt-6.1-sol",
|
||
retry_backoff_seconds=0)
|
||
|
||
# -- 认领门 --------------------------------------------------------
|
||
def test_claim_time_quota_stop(self):
|
||
_insert_quota(self.db_path, "tester", QUOTA)
|
||
_insert_donor_cost(self.db_path, QUOTA)
|
||
_insert_task(self.db_path, "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
|
||
self.calls))
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertIn("quota_exceeded", result["error"])
|
||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||
self.assertIn("quota_exceeded", _error_text(self.db_path, "T1"))
|
||
# 还没跑：runner 一次没调，也不记成本。
|
||
self.assertEqual(self.calls, [])
|
||
self.assertEqual(_cost_rows(self.db_path, "T1"), [])
|
||
|
||
# -- 后续消息执行时查配额 ------------------------------------------
|
||
def test_followup_quota_stop(self):
|
||
_insert_quota(self.db_path, "tester", QUOTA)
|
||
_insert_donor_cost(self.db_path, 90000)
|
||
_insert_task(self.db_path, "T1")
|
||
_insert_followup(self.db_path, "fu1", "T1")
|
||
# 首轮花 30000 微美元：90000 + 30000 = 120000 >= 100000，
|
||
# 排空时必须停，该条回 pending，任务 interrupted。
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 10000, "cached_input_tokens": 0,
|
||
"output_tokens": 1000, "total_tokens": 11000},
|
||
self.calls))
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertIn("quota_exceeded", result["error"])
|
||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||
self.assertEqual(_followup_status(self.db_path, "fu1"), "pending")
|
||
# 首轮只跑了一次（后续消息没跑），已花成本照记 actual。
|
||
self.assertEqual(len(self.calls), 1)
|
||
self.assertEqual(_cost_rows(self.db_path, "T1"),
|
||
[(30000, "actual")])
|
||
|
||
def test_followup_runs_when_under_quota(self):
|
||
_insert_quota(self.db_path, "tester", QUOTA)
|
||
_insert_task(self.db_path, "T1")
|
||
_insert_followup(self.db_path, "fu1", "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 100, "cached_input_tokens": 0,
|
||
"output_tokens": 10, "total_tokens": 110},
|
||
self.calls))
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(_status(self.db_path, "T1"), "completed")
|
||
self.assertEqual(_followup_status(self.db_path, "fu1"), "done")
|
||
self.assertEqual(len(self.calls), 2)
|
||
|
||
def test_no_quota_row_is_unlimited(self):
|
||
_insert_task(self.db_path, "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 100, "output_tokens": 10, "total_tokens": 110},
|
||
self.calls))
|
||
task = worker.get_task("T1")
|
||
self.assertIsNone(worker._quota_exceeded(task, None))
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "completed")
|
||
|
||
# -- 无价 token 上限 -------------------------------------------------
|
||
def test_unpriced_token_cap_blocks_claim(self):
|
||
_insert_quota(self.db_path, "tester", 10 ** 12) # USD 很大，不挡
|
||
_insert_unpriced_cost(self.db_path, 150)
|
||
_insert_task(self.db_path, "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
|
||
self.calls))
|
||
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
|
||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
|
||
try:
|
||
result = worker.run_one()
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
|
||
else:
|
||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertIn("无价模型 token", result["error"])
|
||
self.assertEqual(self.calls, [])
|
||
|
||
def test_unpriced_round_usage_counts_toward_cap(self):
|
||
_insert_quota(self.db_path, "tester", 10 ** 12)
|
||
_insert_unpriced_cost(self.db_path, 90)
|
||
_insert_task(self.db_path, "T1")
|
||
_insert_followup(self.db_path, "fu1", "T1")
|
||
worker = Worker(self.db_path, "/fake/host.py", runner=_usage_runner(
|
||
{"input_tokens": 50, "cached_input_tokens": 0,
|
||
"output_tokens": 0, "total_tokens": 50},
|
||
self.calls), env=self.env, model="gpt-5.6-terra",
|
||
retry_backoff_seconds=0)
|
||
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
|
||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
|
||
try:
|
||
result = worker.run_one()
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
|
||
else:
|
||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
|
||
# 首轮无价 50 token：90 + 50 = 140 >= 100，排空时停。
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertIn("quota_exceeded", result["error"])
|
||
self.assertEqual(_followup_status(self.db_path, "fu1"), "pending")
|
||
|
||
def test_priced_usage_ignores_token_cap(self):
|
||
_insert_quota(self.db_path, "tester", 10 ** 12)
|
||
_insert_unpriced_cost(self.db_path, 90)
|
||
_insert_task(self.db_path, "T1")
|
||
_insert_followup(self.db_path, "fu1", "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 100, "cached_input_tokens": 0,
|
||
"output_tokens": 10, "total_tokens": 110},
|
||
self.calls))
|
||
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
|
||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
|
||
try:
|
||
result = worker.run_one()
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
|
||
else:
|
||
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
|
||
# 有价模型用量不计入无价上限：90 < 100，正常跑完。
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(_followup_status(self.db_path, "fu1"), "done")
|
||
|
||
# -- 轮询掐断 ------------------------------------------------------
|
||
def test_poll_loop_quota_kill(self):
|
||
_insert_task(self.db_path, "T1")
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env=self.env, model="gpt-6.1-sol",
|
||
process_factory=lambda task: subprocess.Popen(
|
||
["sleep", "30"], start_new_session=True,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE))
|
||
task = worker.get_task("T1")
|
||
result = worker._execute_via_process(
|
||
task, kind="dispatch", quota_check=lambda: True)
|
||
self.assertEqual(result["status"], "quota_exceeded")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|