"""P1 执行面针对性单测（只测本路改动，不跑全量三套）。 运行（在 ~/workspace/codex-app-server-p1security 下）： python3 -m unittest discover -s cloud/tests -p "test_p1_security_worker.py" -v 说明：只用注入的假 runner 与测试假值，不跑真实模型回合、不联网、 不登录任何服务；env 里的值只是占位用的假名，绝不是真实密钥。 本文件只归本路（执行面）所有，不碰 test_p1_security_server.py。 """ import json import os import sqlite3 import sys import tempfile import threading import time import unittest from pathlib import Path from types import SimpleNamespace from unittest import mock try: import cloud.sandbox as sandbox_mod from cloud.cost import CostStore from cloud.pr_exec import enqueue_pr_job from cloud.sandbox import DockerExecutor, LocalExecutor from cloud.server import init_db from cloud.worker import Worker, _code_signature except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path sys.path.insert(0, str(Path(__file__).resolve().parents[2])) import cloud.sandbox as sandbox_mod from cloud.cost import CostStore from cloud.pr_exec import enqueue_pr_job from cloud.sandbox import DockerExecutor, LocalExecutor from cloud.server import init_db from cloud.worker import Worker, _code_signature FAKE_ENV = {"XCM_API_KEY": "test-placeholder"} def _completed_runner(cmd): if "dispatch" in cmd: name = cmd[cmd.index("--name") + 1] else: name = cmd[cmd.index("--task") + 1] return SimpleNamespace( returncode=0, stdout=json.dumps({"task": name, "status": "completed", "text": "done"}), stderr="") def _usage_runner(usage, calls): def fake(cmd): calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout=json.dumps({"status": "completed", "text": "done", "usage": usage}), stderr="") return fake def _insert_task(db_path, task_id, created_by="tester", status="queued", project_id="p1", **cols): 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), ) for key, value in cols.items(): conn.execute( "UPDATE tasks SET %s = ? WHERE id = ?" % key, (value, task_id)) conn.commit() finally: conn.close() def _insert_quota(db_path, user_id, micros): conn = sqlite3.connect(db_path) try: conn.execute( "CREATE TABLE IF NOT EXISTS user_quotas(" "user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT," " updated_at DOUBLE PRECISION)") 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_cost(db_path, task_id, model, input_tokens, output_tokens, micros, source="actual"): _insert_task(db_path, task_id, status="completed") store = CostStore(db_path) try: store.record(task_id, "p1", model, input_tokens=input_tokens, output_tokens=output_tokens, usd_micros=micros, source=source) finally: try: store._conn.close() except Exception: pass 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 class P1WorkerTestCase(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.calls = [] def tearDown(self): self.tmp.cleanup() def _worker(self, runner=None, **kwargs): kwargs.setdefault("env", FAKE_ENV) return Worker(self.db_path, "/fake/host.py", runner=runner or _completed_runner, **kwargs) # 1. 无配额行按默认 $10 额度（不再是不限额） def test_missing_quota_row_uses_default_quota(self): _insert_cost(self.db_path, "donor-a", "gpt-6.1-sol", 1, 0, 11000000) _insert_cost(self.db_path, "donor-b", "gpt-6.1-sol", 1, 0, 9000000) worker = self._worker() over = worker._quota_exceeded( {"created_by": "tester", "project_id": "p1"}, None) # donor-a 归 tester：已花 $11 > 默认 $10，挡住。 self.assertIsNotNone(over) self.assertIn("quota_exceeded", over) # donor-b 同样归 tester 会把总额推得更高，单独测另一用户： _insert_task(self.db_path, "donor-c", created_by="rich", status="completed") store = CostStore(self.db_path) try: store.record("donor-c", "p1", "gpt-6.1-sol", input_tokens=1, output_tokens=0, usd_micros=9000000, source="actual") finally: try: store._conn.close() except Exception: pass under = worker._quota_exceeded( {"created_by": "rich", "project_id": "p1"}, None) self.assertIsNone(under) # 2. 显式 NULL 配额行仍是不限额 def test_explicit_null_quota_row_stays_unlimited(self): _insert_quota(self.db_path, "tester", None) _insert_cost(self.db_path, "donor", "gpt-6.1-sol", 1, 0, 999000000) worker = self._worker() self.assertIsNone(worker._quota_exceeded( {"created_by": "tester", "project_id": "p1"}, None)) # 3. 属主为空拒绝新回合（不再返回 None 放行） def test_ownerless_task_refused(self): worker = self._worker() for owner in (None, ""): err = worker._quota_exceeded( {"created_by": owner, "project_id": "p1"}, None) self.assertIsNotNone(err) self.assertIn("quota_exceeded", err) # 4. 协调员子任务属主跟父任务走 def test_child_inherits_parent_owner(self): _insert_quota(self.db_path, "tester", 100) _insert_cost(self.db_path, "donor", "gpt-6.1-sol", 1, 0, 200) _insert_task(self.db_path, "parent", created_by="tester") _insert_task(self.db_path, "parent2", created_by="fresh") worker = self._worker() blocked = worker._quota_exceeded( {"created_by": None, "parent_task_id": "parent", "project_id": "p1"}, None) self.assertIsNotNone(blocked) self.assertIn("quota_exceeded", blocked) # fresh 用户无配额行、未花钱：默认额度内放行。 self.assertIsNone(worker._quota_exceeded( {"created_by": None, "parent_task_id": "parent2", "project_id": "p1"}, None)) # 5. 查库失败 fail-closed（新回合拒绝，不自动放行） def test_quota_db_failure_is_fail_closed(self): worker = self._worker() with mock.patch("cloud.worker._db_connect", side_effect=RuntimeError("test db down")): err = worker._quota_exceeded( {"created_by": "tester", "project_id": "p1"}, None) self.assertIsNotNone(err) self.assertTrue(err.startswith("quota_check_failed:")) # 6. Docker 配额看门狗：命中杀容器、未命中不杀 def test_docker_quota_watchdog(self): killed = [] def fake(cmd, timeout=None): killed.append(list(cmd)) return SimpleNamespace(returncode=0, stdout="", stderr="") executor = DockerExecutor(runner=fake) worker = Worker(self.db_path, "/fake/host.py", env=FAKE_ENV, executor=executor) # 未命中：0.2 秒后停表，kill 不应被调。 stop = threading.Event() threading.Timer(0.2, stop.set).start() worker._docker_quota_watchdog( "T1", "codex-test", stop, lambda: False, interval=0.05) self.assertEqual(killed, []) # 命中：立即杀容器并记命中标记。 stop2 = threading.Event() worker._docker_quota_watchdog( "T2", "codex-test", stop2, lambda: True, interval=0.05) self.assertTrue(any("kill" in c for c in killed)) self.assertTrue(worker._take_quota_hit("T2")) # 命中已消费，第二次取不到（不污染下一轮）。 self.assertFalse(worker._take_quota_hit("T2")) # 7. usage_cost_usd 按实际成本累加，max_cost_usd 不再是死参数 def test_usage_cost_accumulates_and_budget_trips(self): _insert_task(self.db_path, "T1") worker = self._worker( runner=_usage_runner( {"input_tokens": 10000, "cached_input_tokens": 0, "output_tokens": 1000, "total_tokens": 11000}, self.calls), model="gpt-6.1-sol", max_cost_usd=0.01) self.assertFalse(worker._budget_exceeded()) result = worker.run_one() self.assertEqual(result["status"], "completed") # 10000*$2 + 1000*$10 = $0.03，超过 $0.01 上限。 self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6) self.assertTrue(worker._budget_exceeded()) conn = sqlite3.connect(self.db_path) try: row = conn.execute( "SELECT usd_micros, source FROM cost_records" " WHERE task_id = 'T1'").fetchone() finally: conn.close() self.assertEqual(row, (30000, "actual")) # 8. 取消分支先记成本再返回（与配额中断收尾同款记账点） def test_cancel_branch_records_cost(self): worker = self._worker() task = {"id": "T1", "project_id": "p1", "prompt": "do work"} with worker._db_lock: cost_id = worker._record_usage_cost( task, "gpt-6.1-sol", {"input_tokens": 10000, "cached_input_tokens": 0, "output_tokens": 1000, "total_tokens": 11000}, "done", 0) self.assertIsNotNone(cost_id) self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6) # 无用量时回落估算，不累加 usage_cost_usd。 with worker._db_lock: worker._record_usage_cost( {"id": "T2", "project_id": "p1", "prompt": "do work"}, "gpt-6.1-sol", None, "done", 0) self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6) # 9. 估算行不计入无价 token 上限 def test_estimated_rows_excluded_from_unpriced_cap(self): _insert_cost(self.db_path, "donor-est", "gpt-5.6-terra", 100000, 0, None, source="estimated") worker = self._worker() old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP") os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100" try: # 10 万估算 token 在上限 100 下仍放行（估算不是真花费）。 self.assertIsNone(worker._quota_exceeded( {"created_by": "tester", "project_id": "p1"}, None)) _insert_cost(self.db_path, "donor-act", "gpt-5.6-terra", 75, 75, None, source="actual") err = worker._quota_exceeded( {"created_by": "tester", "project_id": "p1"}, None) self.assertIsNotNone(err) self.assertIn("无价模型 token", err) finally: if old is None: os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None) else: os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old # 10. cache_write 不计入无价上限（只 input+output 计） def test_cache_write_excluded_from_unpriced_cap(self): worker = self._worker() old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP") os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100" try: usage = {"input_tokens": 10, "output_tokens": 5, "cache_write_input_tokens": 9999999, "total_tokens": 15} self.assertIsNone(worker._quota_exceeded( {"created_by": "tester", "project_id": "p1", "model": "gpt-5.6-terra"}, usage)) usage2 = {"input_tokens": 200, "output_tokens": 0, "cache_write_input_tokens": 0, "total_tokens": 200} err = worker._quota_exceeded( {"created_by": "tester", "project_id": "p1", "model": "gpt-5.6-terra"}, usage2) self.assertIsNotNone(err) finally: if old is None: os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None) else: os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old # 11. 协调员子任务继承父任务模型/强度，分工显式值优先 def test_coordinator_child_inherits_model_effort(self): worker = self._worker() task = {"id": "coord", "project_id": "p1", "prompt": "coord", "priority": 0, "created_by": "tester", "model": "gpt-6.1-sol", "effort": "high", "executor": None, "environment_id": None} children = worker._coord_sync_children(task, [ {"name": "a", "prompt": "do a", "model": None, "effort": None, "executor": None}, {"name": "b", "prompt": "do b", "model": "gpt-5.6-terra", "effort": "low", "executor": None}, ]) self.assertEqual(len(children), 2) conn = sqlite3.connect(self.db_path) try: rows = {r[0]: (r[1], r[2]) for r in conn.execute( "SELECT coordinator_part, model, effort FROM tasks" " WHERE parent_task_id = 'coord'").fetchall()} finally: conn.close() self.assertEqual(rows["a"], ("gpt-6.1-sol", "high")) self.assertEqual(rows["b"], ("gpt-5.6-terra", "low")) # 12. 断点续跑：Local/Docker/popen 三路都不传 --goal-text def test_goal_resume_breakpoint_omits_goal_text(self): calls = [] def fake(cmd): calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") task = {"id": "T1", "project_id": "p1", "prompt": "orig goal", "verify_command": "true", "goal_max_rounds": 3} local = LocalExecutor("/fake/host.py", runner=fake) local.goal_resume_task(task, None) self.assertIn("goal-resume", calls[-1]) self.assertNotIn("--goal-text", calls[-1]) docker_calls = [] def dfake(cmd, timeout=None): docker_calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") docker = DockerExecutor(runner=dfake, env=FAKE_ENV) with mock.patch("cloud.sandbox.shutil.which", return_value="/usr/bin/docker"): docker.goal_resume_task(task, None) self.assertIn("goal-resume", docker_calls[-1]) self.assertNotIn("--goal-text", docker_calls[-1]) with mock.patch.object(sandbox_mod.subprocess, "Popen") as popen: popen.return_value = SimpleNamespace(pid=9999) local.popen_task( {"id": "T1", "prompt": "resume msg", "_resume_kind": "breakpoint"}) # popen_task 默认 kind=dispatch；显式测 goal-resume： local.popen_task( {"id": "T1", "prompt": "resume msg", "_resume_kind": "breakpoint"}, kind="goal-resume") cmd = popen.call_args[0][0] self.assertIn("goal-resume", cmd) self.assertNotIn("--goal-text", cmd) # 12b. 断点续跑透传 model_provider：首轮 thread 在该 # provider 下创建，续跑命令必须带同一 --model-provider， # 否则 thread/resume 因 provider 表里没有它而失败 # （2026-10-29 probe15 实弹：sub2api not found）。 def test_goal_resume_breakpoint_forwards_model_provider(self): docker_calls = [] def dfake(cmd, timeout=None): docker_calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") task = {"id": "T1", "prompt": "orig goal"} docker = DockerExecutor(runner=dfake, env=FAKE_ENV) with mock.patch("cloud.sandbox.shutil.which", return_value="/usr/bin/docker"): docker.goal_resume_breakpoint_task( task, model_provider="sub2api") cmd = docker_calls[-1] self.assertIn("goal-resume", cmd) self.assertNotIn("--goal-text", cmd) self.assertIn("--model-provider", cmd) self.assertIn("sub2api", cmd) # 不传 provider 时不带该 flag（透传语义，不是强制默认值）。 docker_calls.clear() with mock.patch("cloud.sandbox.shutil.which", return_value="/usr/bin/docker"): docker.goal_resume_breakpoint_task(task) self.assertNotIn("--model-provider", docker_calls[-1]) calls = [] def fake(cmd): calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") local = LocalExecutor("/fake/host.py", runner=fake) local.goal_resume_breakpoint_task(task, model_provider="sub2api") self.assertIn("--model-provider", calls[-1]) self.assertIn("sub2api", calls[-1]) # 12c. worker 本地断点调用点透传网关 provider：run_one # 认领 goal_resume_pending 任务后，本地断点命令必须带 # 网关解析出的 --model-provider（Docker 调用点同理， # 沙箱级已覆盖命令形状）。 def test_worker_local_breakpoint_forwards_gateway_provider(self): def fake(cmd): self.calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") worker = self._worker(runner=fake) worker._ensure_gateway = lambda task: ("sub2api", {}) _insert_task(self.db_path, "T1", goal_mode=1, goal="orig goal", goal_max_rounds=3, goal_resume_pending=1, model="gpt-5.6-terra") worker.run_one() self.assertEqual(_status(self.db_path, "T1"), "completed") self.assertTrue(any( "--model-provider" in cmd and "sub2api" in cmd for cmd in self.calls)) # 12d. Docker 断点续跑必须把密钥环境透传给 runner： # 密钥类变量只按名进 docker run（-e SUB2API_API_KEY 不带值）， # 值靠 _docker_cli_env 放进 CLI 子进程环境；漏传 env 时容器 # 内 key 为空，每轮瞬间报 Missing SUB2API_API_KEY（2026-10-29 # 长任务实弹：断点后 274 轮全失败、无限空转）。 def test_docker_breakpoint_forwards_secret_env_to_runner(self): seen = {} def dfake(cmd, timeout=None, env=None): seen["cmd"] = list(cmd) seen["env"] = env return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") task = {"id": "T1", "prompt": "orig goal"} docker = DockerExecutor(runner=dfake, env=FAKE_ENV) with mock.patch("cloud.sandbox.shutil.which", return_value="/usr/bin/docker"): docker.goal_resume_breakpoint_task( task, model_provider="sub2api", extra_env={"SUB2API_API_KEY": "secret-placeholder", "SUB2API_BASE_URL": "https://example.test"}) # 命令里密钥只按名出现，不带值（不进 argv）。 self.assertIn("SUB2API_API_KEY", seen["cmd"]) self.assertNotIn("SUB2API_API_KEY=secret-placeholder", seen["cmd"]) # runner 的 CLI 环境里必须有值，否则容器内 key 为空。 self.assertIsNotNone(seen["env"]) self.assertEqual(seen["env"].get("SUB2API_API_KEY"), "secret-placeholder") # 13. 开新目标：goal-resume 带 --goal-text（旧行为保留） def test_goal_resume_new_goal_keeps_goal_text(self): calls = [] def fake(cmd): calls.append(list(cmd)) return SimpleNamespace( returncode=0, stdout="goal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") task = {"id": "T1", "project_id": "p1", "prompt": "new goal"} local = LocalExecutor("/fake/host.py", runner=fake) local.goal_resume_task(task, "new goal") self.assertIn("--goal-text", calls[-1]) self.assertIn("new goal", calls[-1]) # 14. goal 轮次上限：sandbox loud fail + worker 认领后 failed def test_goal_max_rounds_cap(self): task = {"id": "T1", "project_id": "p1", "prompt": "g", "verify_command": "true", "goal_max_rounds": 101} local = LocalExecutor("/fake/host.py", runner=_completed_runner) with self.assertRaises(ValueError): local.goal_task(task) _insert_task(self.db_path, "T1", verify_command="true", goal_max_rounds=101) worker = self._worker() result = worker.run_one() self.assertEqual(result["status"], "failed") self.assertIn("上限", result["error"]) self.assertEqual(_status(self.db_path, "T1"), "failed") self.assertEqual(self.calls, []) # 15. PR 非 local：显式 needs_local_worker，行状态持久化， # 任务本身不动，标记过的行仍可被 local worker 认领 def test_pr_non_local_needs_local_worker(self): conn = sqlite3.connect(self.db_path) try: job_id = enqueue_pr_job(conn, "T1", "p1", "test-host") finally: conn.close() executor = DockerExecutor(runner=lambda cmd, timeout=None: SimpleNamespace(returncode=0, stdout="", stderr="")) worker = Worker(self.db_path, "/fake/host.py", env=FAKE_ENV, executor=executor, host_name="test-host") results = worker._run_pr_jobs() self.assertEqual(len(results), 1) self.assertEqual(results[0]["status"], "needs_local_worker") self.assertEqual(results[0]["pr_job_id"], job_id) conn = sqlite3.connect(self.db_path) try: row = conn.execute( "SELECT status FROM pr_jobs WHERE id = ?", (job_id,)).fetchone() self.assertEqual(row[0], "needs_local_worker") # 标记过的行仍可被 local worker 认领，不会被搁置。 from cloud.pr_exec import claim_pr_job claimed = claim_pr_job(conn, "test-host") self.assertIsNotNone(claimed) self.assertEqual(claimed["id"], job_id) finally: conn.close() # 16. 多文件签名 + 签名缺失/已变时认领 fail-closed def test_code_signature_multi_file_and_claim_refusal(self): sig = _code_signature("/fake/host.py") self.assertIsNotNone(sig) self.assertEqual(len(sig), 6) for entry in sig: self.assertEqual(len(entry), 2) _insert_task(self.db_path, "T1") worker = self._worker() worker._code_signature = ("stale-signature",) result = worker.run_one() self.assertFalse(result["dispatched"]) self.assertEqual(result["reason"], "already_claimed") self.assertEqual(_status(self.db_path, "T1"), "queued") worker._code_signature = None result = worker.run_one() self.assertFalse(result["dispatched"]) self.assertEqual(_status(self.db_path, "T1"), "queued") # 17. goal 首轮撞已存在 → 断点续跑一次，不直接 failed def test_goal_already_exists_falls_back_to_breakpoint(self): calls = [] def fake(cmd): calls.append(list(cmd)) if "goal-resume" in cmd: return SimpleNamespace( returncode=0, stdout="round 1: turn=completed verify_rc=0 " "verify_out=\ngoal met after rounds\n" "usage input=10 cached=0 cache_write=0 output=5 " "reasoning=0 total=15 completions=1\n", stderr="") if "goal" in cmd: return SimpleNamespace( returncode=0, stdout="host 任务已存在：T1 已存在", stderr="") return _completed_runner(cmd) _insert_task(self.db_path, "T1", verify_command="true", goal_max_rounds=2) worker = self._worker(runner=fake) result = worker.run_one() self.assertEqual(result["status"], "completed") # 注意：cmd 是元素列表，"goal" in cmd 只匹配恰好等于 # "goal" 的元素，goal-resume 命令里并没有 "goal" 元素， # 所以按“任一元素含 goal 子串”过滤。 goal_cmds = [c for c in calls if any("goal" in e for e in c)] self.assertTrue(any("goal-resume" in c for c in goal_cmds)) resume_cmd = [c for c in goal_cmds if "goal-resume" in c][0] self.assertNotIn("--goal-text", resume_cmd) self.assertEqual(_status(self.db_path, "T1"), "completed") if __name__ == "__main__": unittest.main()