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