host
9f49b97cd3
- test_system_quota_exempt_server → test_system_quota_not_exempt_server：system 按普通属主算，quota 0 时照常拦（402）。 - test_system_task_exempt_from_quota → test_system_task_not_exempt_from_quota：system quota 0 时认领门拦，标 interrupted。 - test_poll_quota_lookup_failure_open → test_poll_quota_lookup_failure_open_returns_check_failed：open 异常时回 quota_check_failed（边界拒绝、轮询经 _quota_hit 过滤只警告），closed 回 quota_check_unavailable。
488 lines
20 KiB
Python
488 lines
20 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 threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud.cost import CostStore
|
||
from cloud.sandbox import DockerExecutor, LocalExecutor
|
||
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.sandbox import DockerExecutor, LocalExecutor
|
||
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")
|
||
|
||
# -- G7 分层：门侧 fail-closed、轮询侧 fail-open -------------------
|
||
def test_claim_gate_quota_lookup_failure_closed(self):
|
||
_insert_task(self.db_path, "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||
self.calls))
|
||
with mock.patch.object(
|
||
Worker, "_quota_limit_and_used",
|
||
side_effect=RuntimeError("db down")):
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertIn("quota_check_unavailable", result["error"])
|
||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||
# 门侧拦住：runner 一次没调，也不记成本（cost_records
|
||
# 是 CostStore 懒建的，没跑过就根本没有这张表）。
|
||
self.assertEqual(self.calls, [])
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
has_table = conn.execute(
|
||
"SELECT 1 FROM sqlite_master WHERE type='table'"
|
||
" AND name='cost_records'").fetchone() is not None
|
||
rows = _cost_rows(self.db_path, "T1") if has_table else []
|
||
finally:
|
||
conn.close()
|
||
self.assertEqual(rows, [])
|
||
|
||
def test_poll_quota_lookup_failure_open_returns_check_failed(self):
|
||
# 合并后口径（G7 + P1）：open 不再回 None 放行，而是回
|
||
# quota_check_failed——新回合边界见此串即拒绝；执行中
|
||
# 轮询经 _quota_hit 闭包过滤此串、只警告不杀轮（轮询
|
||
# 不杀的语义由闭包保证，不是由这里回 None 保证）。
|
||
# closed（认领门）回 quota_check_unavailable。
|
||
_insert_task(self.db_path, "T1")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||
self.calls))
|
||
task = worker.get_task("T1")
|
||
with mock.patch.object(
|
||
Worker, "_quota_limit_and_used",
|
||
side_effect=RuntimeError("db down")):
|
||
err = worker._quota_exceeded(task, None)
|
||
self.assertTrue(err.startswith("quota_check_failed:"))
|
||
err = worker._quota_exceeded(task, None, on_error="open")
|
||
self.assertTrue(err.startswith("quota_check_failed:"))
|
||
err = worker._quota_exceeded(task, None, on_error="closed")
|
||
self.assertTrue(err.startswith("quota_check_unavailable:"))
|
||
|
||
# -- G13：system 不豁免（worker 侧，P1 口径） ------------------------
|
||
def test_system_task_not_exempt_from_quota(self):
|
||
# P1 口径：system 按普通属主算，不豁免（G14 的 system 豁免
|
||
# 与 P1 冲突，rebase 时未合入）。quota 0 时认领门即拦，
|
||
# 标 interrupted（可续跑），不是 completed。
|
||
_insert_quota(self.db_path, "system", 0)
|
||
_insert_task(self.db_path, "T1", created_by="system")
|
||
worker = self._worker(_usage_runner(
|
||
{"input_tokens": 100, "output_tokens": 10, "total_tokens": 110},
|
||
self.calls))
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertEqual(_status(self.db_path, "T1"), "interrupted")
|
||
|
||
# -- G8 Docker 轮询掐断 ---------------------------------------------
|
||
def test_docker_quota_poll_kills_container(self):
|
||
_insert_quota(self.db_path, "tester", QUOTA)
|
||
_insert_task(self.db_path, "T1")
|
||
_insert_followup(self.db_path, "fu1", "T1")
|
||
executor = _BlockingDockerExecutor()
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env=self.env, model="gpt-6.1-sol",
|
||
executor=executor, retry_backoff_seconds=0,
|
||
quota_poll_interval_seconds=0.01)
|
||
calls = {"n": 0}
|
||
|
||
def fake_limit(owner):
|
||
calls["n"] += 1
|
||
# 第 1 次是认领门（未超限，放行）；之后轮询一律超限。
|
||
# helper 在 rebase 合并 #86 无价上限后返回 4 元组
|
||
# （quota, used, month, unpriced_used），此处无价
|
||
# 用量为 0，不触发无价上限。
|
||
return (QUOTA, 90000 if calls["n"] == 1 else QUOTA,
|
||
"2026-10", 0)
|
||
|
||
worker._quota_limit_and_used = fake_limit
|
||
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(executor.killed, ["codex-T1"])
|
||
# 续跑消息保留（仍 pending），已花成本照记。
|
||
self.assertEqual(_followup_status(self.db_path, "fu1"), "pending")
|
||
self.assertEqual(len(_cost_rows(self.db_path, "T1")), 1)
|
||
|
||
# -- G9：Worker 把 ignore 开关同步给外部执行器 ----------------------
|
||
def test_worker_syncs_ignore_host_budget_to_executor(self):
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env=self.env, model="gpt-6.1-sol",
|
||
executor=LocalExecutor("/fake/host.py"),
|
||
retry_backoff_seconds=0)
|
||
self.assertTrue(worker.ignore_host_budget)
|
||
self.assertTrue(worker.executor.ignore_host_budget)
|
||
worker2 = Worker(self.db_path, "/fake/host.py",
|
||
env=self.env, model="gpt-6.1-sol",
|
||
executor=LocalExecutor("/fake/host.py"),
|
||
ignore_host_budget=False,
|
||
retry_backoff_seconds=0)
|
||
self.assertFalse(worker2.ignore_host_budget)
|
||
self.assertFalse(worker2.executor.ignore_host_budget)
|
||
|
||
def test_docker_quota_poll_failure_open(self):
|
||
_insert_task(self.db_path, "T1")
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env=self.env, model="gpt-6.1-sol",
|
||
quota_poll_interval_seconds=0.01)
|
||
task = worker.get_task("T1")
|
||
|
||
def slow_fn():
|
||
time.sleep(0.05)
|
||
return {"status": "completed"}
|
||
|
||
def bad_check():
|
||
raise RuntimeError("db down")
|
||
|
||
result = worker._docker_call_with_quota_poll(
|
||
task, slow_fn, bad_check)
|
||
self.assertEqual(result["status"], "completed")
|
||
|
||
|
||
class _BlockingDockerExecutor(DockerExecutor):
|
||
"""G8 单测用假 Docker：dispatch 阻塞到被 kill，绝不调真 docker。"""
|
||
|
||
def __init__(self):
|
||
super().__init__(host_py_path="/fake/host.py")
|
||
self.killed = []
|
||
self._stop = threading.Event()
|
||
|
||
def dispatch_task(self, task, timeout=None, extra_mounts=None,
|
||
network=None, extra_env=None, seed_host_dir=None,
|
||
model_provider=None):
|
||
self._stop.wait(timeout=10)
|
||
return {"status": "failed", "stdout": "", "stderr": "killed",
|
||
"usage": {"input_tokens": 100, "cached_input_tokens": 0,
|
||
"output_tokens": 10, "total_tokens": 110}}
|
||
|
||
def kill_container(self, name):
|
||
self.killed.append(name)
|
||
self._stop.set()
|
||
return True
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|