host
1402dc67fd
- Docker quota watchdog + fail-closed quota checks, system tasks need real owner - max_cost_usd actually accumulated, cancel records cost, estimated rows excluded from unpriced cap, cache_write declared non-billed - goal-resume keeps original verify on breakpoint resume, goal retry falls back to resume, goal_max_rounds capped at 100 at create and execution - coordinator children inherit parent model/effort, orphans marked interrupted with audit - PR jobs on docker workers return explicit needs_local_worker instead of silent empty - worker code signature covers 6 files and is checked before claim, fail-closed - sub2api lazy provisioning with rotation, backup whitelist + 0600, push SSRF block, repos/workers desensitized, key/project audits and limits, console buttons/quota/status wording, coordinator/automation quota reserve
691 lines
29 KiB
Python
691 lines
29 KiB
Python
"""cloud worker 单测：认领顺序、取消、缺 key、失败写回、预算门。
|
||
|
||
运行（在 ~/workspace/codex-app-server 下）：
|
||
python3 -m unittest discover -s cloud/tests -p "test_worker.py" -v
|
||
|
||
说明：只用注入的假 runner（假 CLI），不跑真实模型回合、不联网、
|
||
不登录任何服务；env 映射也是注入的普通 dict，里面的值只是
|
||
占位用的假名，绝不是真实密钥。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
|
||
try:
|
||
from cloud.cost import CostStore
|
||
from cloud.models import ModelRouter
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker, resolve_env_key
|
||
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.models import ModelRouter
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker, resolve_env_key
|
||
|
||
|
||
def _insert_task(db_path, task_id, prompt="do work", status="queued",
|
||
depends_on=None, priority=0, 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, ?, ?, 'tester', ?, ?, NULL)",
|
||
(task_id, project_id, prompt, status,
|
||
json.dumps(depends_on or []), priority, 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 _completed_runner(cmd):
|
||
"""假 CLI：dispatch/run 一律成功，打印 host.py 风格的 JSON。"""
|
||
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="")
|
||
|
||
|
||
class WorkerTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
# P1 SSRF 门默认挡 loopback；通知测试用本地监听器，显式放行，
|
||
# 生产默认值仍关闭。
|
||
self._old_push_allow = os.environ.get(
|
||
"CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
|
||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
|
||
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):
|
||
if getattr(self, "_old_push_allow", None) is None:
|
||
os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
|
||
else:
|
||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = (
|
||
self._old_push_allow)
|
||
self.tmp.cleanup()
|
||
|
||
def _worker(self, runner=None, **kwargs):
|
||
kwargs.setdefault("env", self.env)
|
||
return Worker(self.db_path, "/fake/host.py",
|
||
runner=runner or _completed_runner, **kwargs)
|
||
|
||
def _recording_runner(self, inner):
|
||
def fake(cmd):
|
||
self.calls.append(list(cmd))
|
||
return inner(cmd)
|
||
return fake
|
||
|
||
# -- 认领顺序：依赖拓扑 + 优先级 ------------------------------------
|
||
def test_claim_order_dependency_and_priority(self):
|
||
_insert_task(self.db_path, "A", priority=0)
|
||
_insert_task(self.db_path, "B", depends_on=["A"], priority=0)
|
||
_insert_task(self.db_path, "C", priority=9)
|
||
worker = self._worker(runner=self._recording_runner(_completed_runner))
|
||
results = worker.run_all()
|
||
self.assertEqual([r["status"] for r in results],
|
||
["completed"] * 3)
|
||
# 首波 {C, A} 按优先级降序：C 先于 A；B 依赖 A，只能最后。
|
||
dispatch_names = [c[c.index("--name") + 1] for c in self.calls
|
||
if "dispatch" in c]
|
||
self.assertEqual(dispatch_names, ["C", "A", "B"])
|
||
for task_id in ("A", "B", "C"):
|
||
self.assertEqual(_status(self.db_path, task_id), "completed")
|
||
|
||
def test_dependency_on_failed_task_not_dispatched(self):
|
||
def fake(cmd):
|
||
if "dispatch" in cmd and cmd[cmd.index("--name") + 1] == "A":
|
||
return SimpleNamespace(returncode=1, stdout="boom",
|
||
stderr="")
|
||
return _completed_runner(cmd)
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B", depends_on=["A"])
|
||
worker = self._worker(runner=fake)
|
||
results = worker.run_all()
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["status"], "failed")
|
||
self.assertEqual(_status(self.db_path, "A"), "failed")
|
||
# B 的依赖未完成，始终不可派，保持 queued。
|
||
self.assertEqual(_status(self.db_path, "B"), "queued")
|
||
|
||
# -- 终态推送 --------------------------------------------------------
|
||
def test_notify_on_completed_and_failed(self):
|
||
"""任务终态自动通知订阅：completed 与 failed 各发一次真 POST
|
||
到本地回环 webhook，body 含事件类型与任务 id；取消
|
||
（interrupted）不在事件集合里，不通知（见下个测试）。"""
|
||
import http.server
|
||
import threading
|
||
|
||
requests = []
|
||
|
||
class Handler(http.server.BaseHTTPRequestHandler):
|
||
def do_POST(self):
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
body = self.rfile.read(length) if length else b""
|
||
requests.append(json.loads(body.decode("utf-8")))
|
||
self.send_response(200)
|
||
self.send_header("Content-Length", "2")
|
||
self.end_headers()
|
||
self.wfile.write(b"{}")
|
||
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
server = http.server.ThreadingHTTPServer(
|
||
("127.0.0.1", 0), Handler)
|
||
thread = threading.Thread(
|
||
target=server.serve_forever, daemon=True)
|
||
thread.start()
|
||
try:
|
||
from cloud.push import create_subscription
|
||
url = "http://127.0.0.1:%d/hook" % server.server_address[1]
|
||
create_subscription(
|
||
self.db_path, "p1", "user-1", "webhook", url, "instant")
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B")
|
||
|
||
def failing(cmd):
|
||
if "dispatch" in cmd and cmd[cmd.index("--name") + 1] == "B":
|
||
return SimpleNamespace(returncode=1, stdout="boom",
|
||
stderr="")
|
||
return _completed_runner(cmd)
|
||
|
||
worker = self._worker(runner=failing)
|
||
results = worker.run_all()
|
||
self.assertEqual(
|
||
sorted(r["status"] for r in results),
|
||
["completed", "failed"])
|
||
by_task = {b["task_id"]: b for b in requests}
|
||
self.assertEqual(set(by_task), {"A", "B"})
|
||
self.assertEqual(by_task["A"]["event_type"], "completed")
|
||
self.assertEqual(by_task["B"]["event_type"], "failed")
|
||
self.assertIn("boom", json.dumps(by_task["B"]["payload"]))
|
||
finally:
|
||
server.shutdown()
|
||
server.server_close()
|
||
thread.join(timeout=5)
|
||
|
||
def test_cancel_does_not_notify(self):
|
||
"""取消只标 interrupted，不在推送事件集合里，不发通知。"""
|
||
import http.server
|
||
import threading
|
||
|
||
requests = []
|
||
|
||
class Handler(http.server.BaseHTTPRequestHandler):
|
||
def do_POST(self):
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
self.rfile.read(length) if length else b""
|
||
requests.append(True)
|
||
self.send_response(200)
|
||
self.send_header("Content-Length", "2")
|
||
self.end_headers()
|
||
self.wfile.write(b"{}")
|
||
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
server = http.server.ThreadingHTTPServer(
|
||
("127.0.0.1", 0), Handler)
|
||
thread = threading.Thread(
|
||
target=server.serve_forever, daemon=True)
|
||
thread.start()
|
||
try:
|
||
from cloud.push import create_subscription
|
||
url = "http://127.0.0.1:%d/hook" % server.server_address[1]
|
||
create_subscription(
|
||
self.db_path, "p1", "user-1", "webhook", url, "instant")
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(runner=_completed_runner)
|
||
self.assertTrue(worker.cancel("A"))
|
||
self.assertEqual(worker.run_all(), [])
|
||
self.assertEqual(requests, [])
|
||
finally:
|
||
server.shutdown()
|
||
server.server_close()
|
||
thread.join(timeout=5)
|
||
|
||
# -- goal 模式 -------------------------------------------------------
|
||
def _set_goal(self, task_id, verify="test -f x", goal=None,
|
||
rounds=None):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"UPDATE tasks SET verify_command = ?, goal = ?,"
|
||
" goal_max_rounds = ? WHERE id = ?",
|
||
(verify, goal, rounds, task_id))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
def test_goal_met_notifies_completed(self):
|
||
"""goal 达标：任务 completed，发 completed（不发 goal-unmet）。"""
|
||
import http.server
|
||
import threading
|
||
|
||
requests = []
|
||
|
||
class Handler(http.server.BaseHTTPRequestHandler):
|
||
def do_POST(self):
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
body = self.rfile.read(length) if length else b""
|
||
requests.append(json.loads(body.decode("utf-8")))
|
||
self.send_response(200)
|
||
self.send_header("Content-Length", "2")
|
||
self.end_headers()
|
||
self.wfile.write(b"{}")
|
||
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
server = http.server.ThreadingHTTPServer(
|
||
("127.0.0.1", 0), Handler)
|
||
thread = threading.Thread(
|
||
target=server.serve_forever, daemon=True)
|
||
thread.start()
|
||
try:
|
||
from cloud.push import create_subscription
|
||
url = "http://127.0.0.1:%d/hook" % server.server_address[1]
|
||
create_subscription(
|
||
self.db_path, "p1", "user-1", "webhook", url, "instant")
|
||
_insert_task(self.db_path, "A")
|
||
self._set_goal("A")
|
||
calls = []
|
||
|
||
def goal_runner(cmd):
|
||
calls.append(list(cmd))
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout="round 1: turn=completed verify_rc=0 "
|
||
"verify_out=ok\ngoal met after rounds\n",
|
||
stderr="")
|
||
|
||
worker = self._worker(runner=goal_runner)
|
||
results = worker.run_all()
|
||
self.assertEqual(
|
||
[r["status"] for r in results], ["completed"])
|
||
self.assertIn("goal", calls[0])
|
||
self.assertEqual(len(requests), 1)
|
||
self.assertEqual(requests[0]["event_type"], "completed")
|
||
finally:
|
||
server.shutdown()
|
||
server.server_close()
|
||
thread.join(timeout=5)
|
||
|
||
def test_goal_unmet_notifies_only_goal_unmet(self):
|
||
"""goal 未达：任务 failed，只发 goal-unmet，不重复发 failed。"""
|
||
import http.server
|
||
import threading
|
||
|
||
requests = []
|
||
|
||
class Handler(http.server.BaseHTTPRequestHandler):
|
||
def do_POST(self):
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
body = self.rfile.read(length) if length else b""
|
||
requests.append(json.loads(body.decode("utf-8")))
|
||
self.send_response(200)
|
||
self.send_header("Content-Length", "2")
|
||
self.end_headers()
|
||
self.wfile.write(b"{}")
|
||
|
||
def log_message(self, *args):
|
||
pass
|
||
|
||
server = http.server.ThreadingHTTPServer(
|
||
("127.0.0.1", 0), Handler)
|
||
thread = threading.Thread(
|
||
target=server.serve_forever, daemon=True)
|
||
thread.start()
|
||
try:
|
||
from cloud.push import create_subscription
|
||
url = "http://127.0.0.1:%d/hook" % server.server_address[1]
|
||
create_subscription(
|
||
self.db_path, "p1", "user-1", "webhook", url, "instant")
|
||
_insert_task(self.db_path, "A")
|
||
self._set_goal("A")
|
||
|
||
def unmet_runner(cmd):
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout="round 1: turn=completed verify_rc=1 "
|
||
"verify_out=file missing\n"
|
||
"goal unmet after rounds\n",
|
||
stderr="")
|
||
|
||
worker = self._worker(runner=unmet_runner)
|
||
results = worker.run_all()
|
||
self.assertEqual([r["status"] for r in results], ["failed"])
|
||
self.assertEqual(len(requests), 1)
|
||
self.assertEqual(requests[0]["event_type"], "goal-unmet")
|
||
self.assertIn("file missing",
|
||
requests[0]["payload"]["error"])
|
||
finally:
|
||
server.shutdown()
|
||
server.server_close()
|
||
thread.join(timeout=5)
|
||
|
||
def test_goal_mode_runs_once_despite_retries(self):
|
||
"""goal 模式只跑一次：循环内部已多轮重试，worker 级重试
|
||
不再叠加（同任务重跑会撞 host 任务名已存在）。"""
|
||
_insert_task(self.db_path, "A")
|
||
self._set_goal("A")
|
||
calls = []
|
||
|
||
def counting(cmd):
|
||
calls.append(cmd)
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout="round 1: turn=completed verify_rc=1 "
|
||
"verify_out=nope\ngoal unmet after rounds\n",
|
||
stderr="")
|
||
|
||
worker = self._worker(runner=counting, max_retries=2)
|
||
results = worker.run_all()
|
||
self.assertEqual(len(calls), 1)
|
||
self.assertEqual(results[0]["status"], "failed")
|
||
|
||
def test_cloud_executor_goal_unsupported(self):
|
||
"""cloud 执行器没有 goal 概念：goal 任务如实失败，
|
||
不静默降级成普通 dispatch。"""
|
||
from cloud.sandbox import CloudExecutor
|
||
_insert_task(self.db_path, "A")
|
||
self._set_goal("A")
|
||
worker = self._worker(executor=CloudExecutor())
|
||
results = worker.run_all()
|
||
self.assertEqual(results[0]["status"], "failed")
|
||
self.assertIn("不支持 goal", results[0]["error"])
|
||
|
||
# -- 成本：真实用量 --------------------------------------------------
|
||
def test_cost_actual_from_usage(self):
|
||
"""执行器带回用量时记 actual：token 全是实测值，usd 按
|
||
价目表算（(1000-400)*2 + 400*0.10 + 100*10 = 2240 微美元）。"""
|
||
_insert_task(self.db_path, "A")
|
||
usage = {"input_tokens": 1000, "cached_input_tokens": 400,
|
||
"cache_write_input_tokens": 0, "output_tokens": 100,
|
||
"reasoning_output_tokens": 20, "total_tokens": 1100,
|
||
"completions": 1}
|
||
|
||
def runner(cmd):
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout=json.dumps({"status": "completed", "text": "done",
|
||
"usage": usage}),
|
||
stderr="")
|
||
|
||
worker = self._worker(runner=runner, model="gpt-6.1-sol")
|
||
results = worker.run_all()
|
||
self.assertEqual([r["status"] for r in results], ["completed"])
|
||
from cloud.cost import CostStore
|
||
report = CostStore(self.db_path).report_by_task("A")
|
||
self.assertEqual(report["actual_tokens"], 1100)
|
||
self.assertEqual(report["cached_input_tokens"], 400)
|
||
self.assertEqual(report["estimated_tokens"], 0)
|
||
self.assertEqual(report["usd_micros"], 2240)
|
||
self.assertEqual(report["unpriced_actual_tokens"], 0)
|
||
|
||
def test_cost_actual_unpriced_model(self):
|
||
"""无价模型：token 照样记 actual，usd 记 None、报表里
|
||
进 unpriced_actual_tokens，绝不按 0 美元处理。"""
|
||
_insert_task(self.db_path, "A")
|
||
usage = {"input_tokens": 1000, "cached_input_tokens": 0,
|
||
"cache_write_input_tokens": 0, "output_tokens": 100,
|
||
"reasoning_output_tokens": 0, "total_tokens": 1100,
|
||
"completions": 1}
|
||
|
||
def runner(cmd):
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout=json.dumps({"status": "completed", "text": "done",
|
||
"usage": usage}),
|
||
stderr="")
|
||
|
||
worker = self._worker(runner=runner, model="gpt-5.6-sol")
|
||
results = worker.run_all()
|
||
self.assertEqual([r["status"] for r in results], ["completed"])
|
||
from cloud.cost import CostStore
|
||
report = CostStore(self.db_path).report_by_task("A")
|
||
self.assertEqual(report["actual_tokens"], 1100)
|
||
self.assertEqual(report["usd_micros"], 0)
|
||
self.assertEqual(report["unpriced_actual_tokens"], 1100)
|
||
|
||
def test_cost_estimated_when_no_usage(self):
|
||
"""执行器没带回用量时回落字符估算记 estimated（旧行为）。"""
|
||
_insert_task(self.db_path, "A")
|
||
|
||
def runner(cmd):
|
||
return SimpleNamespace(returncode=0, stdout="plain text",
|
||
stderr="")
|
||
|
||
worker = self._worker(runner=runner, model="gpt-6.1-sol")
|
||
results = worker.run_all()
|
||
self.assertEqual([r["status"] for r in results], ["completed"])
|
||
from cloud.cost import CostStore
|
||
report = CostStore(self.db_path).report_by_task("A")
|
||
self.assertEqual(report["actual_tokens"], 0)
|
||
self.assertGreater(report["estimated_tokens"], 0)
|
||
|
||
# -- 取消 ------------------------------------------------------------
|
||
def test_cancel_queued_task_marks_interrupted(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(runner=self._recording_runner(
|
||
_completed_runner))
|
||
self.assertTrue(worker.cancel("A"))
|
||
self.assertEqual(_status(self.db_path, "A"), "interrupted")
|
||
# 已终态：再次取消返回 False；不存在的任务也是 False。
|
||
self.assertFalse(worker.cancel("A"))
|
||
self.assertFalse(worker.cancel("NOPE"))
|
||
# 已取消的任务不再被派发。
|
||
self.assertEqual(worker.run_all(), [])
|
||
self.assertEqual(self.calls, [])
|
||
|
||
def test_cancel_during_dispatch_keeps_interrupted(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = None
|
||
|
||
def fake(cmd):
|
||
# 执行中途从外部取消：worker 返回后不得用 completed 覆盖。
|
||
worker.cancel("A")
|
||
return _completed_runner(cmd)
|
||
|
||
worker = self._worker(runner=fake)
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "interrupted")
|
||
self.assertEqual(_status(self.db_path, "A"), "interrupted")
|
||
|
||
# -- 缺 key 门 --------------------------------------------------------
|
||
def test_missing_key_fails_clearly_and_does_not_run(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(
|
||
runner=self._recording_runner(_completed_runner), env={})
|
||
result = worker.run_one()
|
||
self.assertTrue(result["dispatched"])
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertEqual(result["error"],
|
||
"Missing environment variable: XCM_API_KEY")
|
||
self.assertEqual(_status(self.db_path, "A"), "failed")
|
||
self.assertEqual(self.calls, []) # 缺 key 时绝不调 host.py
|
||
|
||
def test_missing_key_empty_value_counts_as_missing(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(env={"XCM_API_KEY": ""})
|
||
result = worker.run_one()
|
||
self.assertEqual(result["error"],
|
||
"Missing environment variable: XCM_API_KEY")
|
||
|
||
# -- 提前失败路的暂存清理 --------------------------------------------
|
||
# 2026-10-27 实弹：仓库解析失败的任务不走正常终局清理，
|
||
# bundle 目录永久残留。缺 key 门是同款提前返回，同测。
|
||
def _attach_bundle(self, task_id, bundle_id="bundle1"):
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO task_attachments(task_id, bundle_id,"
|
||
" rel_path, size) VALUES(?, ?, ?, ?)",
|
||
(task_id, bundle_id, "a.txt", 3))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
bundle_dir = os.path.join(self.upload_dir, bundle_id)
|
||
os.makedirs(bundle_dir, exist_ok=True)
|
||
with open(os.path.join(bundle_dir, "a.txt"), "w") as f:
|
||
f.write("x")
|
||
return bundle_dir
|
||
|
||
def test_repo_failure_cleans_staging_bundle(self):
|
||
_insert_task(self.db_path, "A")
|
||
self.upload_dir = os.path.join(self.tmp.name, "uploads")
|
||
bundle_dir = self._attach_bundle("A")
|
||
worker = self._worker(
|
||
runner=self._recording_runner(_completed_runner),
|
||
upload_dir=self.upload_dir)
|
||
worker._project_repo_url = lambda pid: "https://example.invalid/r.git"
|
||
|
||
def _boom(cmd, timeout=None):
|
||
raise Exception("clone boom")
|
||
|
||
worker.repos_runner = _boom
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertIn("仓库解析失败", result["error"])
|
||
self.assertEqual(_status(self.db_path, "A"), "failed")
|
||
self.assertEqual(self.calls, []) # 解析失败时绝不调 host.py
|
||
self.assertFalse(os.path.exists(bundle_dir))
|
||
|
||
def test_missing_key_cleans_staging_bundle(self):
|
||
_insert_task(self.db_path, "A")
|
||
self.upload_dir = os.path.join(self.tmp.name, "uploads")
|
||
bundle_dir = self._attach_bundle("A")
|
||
worker = self._worker(env={}, upload_dir=self.upload_dir)
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertFalse(os.path.exists(bundle_dir))
|
||
|
||
# -- 代码变化自退 ------------------------------------------------------
|
||
# 2026-10-27 实弹：05:51 的旧 docker worker 成孤儿后继续用旧
|
||
# 代码抢任务。常驻循环每轮边界比对 worker.py 签名，变了就退
|
||
# 出（75，失败类，supervisor 重拉新版），绝不中断在跑任务。
|
||
def test_code_changed_exits_loop_before_claiming(self):
|
||
worker = self._worker(
|
||
runner=self._recording_runner(_completed_runner))
|
||
worker._code_signature = ("stale-signature",)
|
||
with self.assertRaises(SystemExit) as ctx:
|
||
worker.run_loop(interval=10, max_seconds=30)
|
||
self.assertEqual(ctx.exception.code, 75)
|
||
self.assertEqual(self.calls, []) # 自退在认领之前
|
||
|
||
def test_code_unchanged_loop_runs_to_max_seconds(self):
|
||
worker = self._worker()
|
||
results = worker.run_loop(interval=10, max_seconds=0.4)
|
||
self.assertEqual(results, [])
|
||
|
||
def test_custom_env_key_name_from_arg_and_config(self):
|
||
# 参数指定的名字出现在报错里（只含名字，不含值）。
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(env={}, env_key="MY_RELAY_KEY")
|
||
self.assertEqual(worker.run_one()["error"],
|
||
"Missing environment variable: MY_RELAY_KEY")
|
||
# codex 配置文件里的 env_key 行只用来取名字。
|
||
cfg = os.path.join(self.tmp.name, "config.toml")
|
||
with open(cfg, "w", encoding="utf-8") as f:
|
||
f.write('env_key = "CUSTOM_RELAY_KEY"\n')
|
||
self.assertEqual(resolve_env_key(codex_config_path=cfg),
|
||
"CUSTOM_RELAY_KEY")
|
||
_insert_task(self.db_path, "B")
|
||
worker2 = self._worker(env={}, codex_config_path=cfg)
|
||
self.assertEqual(worker2.env_key, "CUSTOM_RELAY_KEY")
|
||
|
||
def test_key_present_proceeds(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker()
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertIsNone(result["error"])
|
||
|
||
# -- 执行失败写回 + 成本 ----------------------------------------------
|
||
def test_execution_failure_writes_back_failed_and_cost(self):
|
||
def fake(cmd):
|
||
return SimpleNamespace(returncode=1, stdout="boom", stderr="")
|
||
_insert_task(self.db_path, "A", prompt="do work")
|
||
worker = self._worker(runner=fake)
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertIn("boom", result["error"])
|
||
self.assertEqual(_status(self.db_path, "A"), "failed")
|
||
store = CostStore(self.db_path)
|
||
try:
|
||
report = store.report_by_task("A")
|
||
row = store._conn.execute(
|
||
"SELECT model, source, fallback_count FROM cost_records"
|
||
" WHERE task_id = ?", ("A",)).fetchone()
|
||
finally:
|
||
store._conn.close()
|
||
self.assertEqual(report["turns"], 1)
|
||
self.assertGreater(report["estimated_tokens"], 0)
|
||
self.assertEqual(report["actual_tokens"], 0)
|
||
# 复用 cost.py 的记录结构：估算标 estimated，不冒充 actual。
|
||
# 模型列跟 worker 默认模型走（2026-10-27 起为 terra）。
|
||
self.assertEqual(row[0], "gpt-5.6-terra")
|
||
self.assertEqual(row[1], "estimated")
|
||
self.assertEqual(row[2], 0)
|
||
|
||
# -- 预算门 ------------------------------------------------------------
|
||
def test_budget_exceeded_no_dispatch(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(
|
||
runner=self._recording_runner(_completed_runner),
|
||
max_turns=0)
|
||
result = worker.run_one()
|
||
self.assertFalse(result["dispatched"])
|
||
self.assertEqual(result["reason"], "budget_exceeded")
|
||
self.assertEqual(_status(self.db_path, "A"), "queued")
|
||
self.assertEqual(self.calls, [])
|
||
|
||
def test_cost_budget_exceeded_no_dispatch(self):
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(max_cost_usd=0.0)
|
||
result = worker.run_one()
|
||
self.assertFalse(result["dispatched"])
|
||
self.assertEqual(_status(self.db_path, "A"), "queued")
|
||
|
||
# -- fallback 事件：只记、不自动重试 -----------------------------------
|
||
def test_timeout_fallback_event_recorded_not_retried(self):
|
||
def fake(cmd):
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout=json.dumps({"task": "A", "status": "timeout",
|
||
"text": ""}),
|
||
stderr="")
|
||
_insert_task(self.db_path, "A")
|
||
router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
|
||
# 显式指定主模型：fallback 场景不依赖全局默认模型的值。
|
||
worker = self._worker(runner=self._recording_runner(fake),
|
||
router=router, model="gpt-6.1-sol")
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertEqual(len(result["events"]), 1)
|
||
event = result["events"][0]
|
||
self.assertEqual(event["type"], "model_fallback")
|
||
self.assertEqual(event["from_model"], "gpt-6.1-sol")
|
||
self.assertEqual(event["to_model"], "gpt-6-sol")
|
||
self.assertEqual(event["reason"], "timeout")
|
||
# 只记事件不重试：假 CLI 只被调用一次。
|
||
self.assertEqual(len(self.calls), 1)
|
||
store = CostStore(self.db_path)
|
||
try:
|
||
report = store.report_by_task("A")
|
||
finally:
|
||
store._conn.close()
|
||
self.assertEqual(report["fallback_count"], 1)
|
||
|
||
# -- dispatch 已存在时用 run 续跑 --------------------------------------
|
||
def test_dispatch_existing_falls_back_to_run(self):
|
||
def fake(cmd):
|
||
if "dispatch" in cmd:
|
||
return SimpleNamespace(
|
||
returncode=1, stdout="",
|
||
stderr="Traceback: task 已存在：A，先换名或 discard 它")
|
||
return _completed_runner(cmd)
|
||
_insert_task(self.db_path, "A")
|
||
worker = self._worker(runner=self._recording_runner(fake))
|
||
result = worker.run_one()
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(_status(self.db_path, "A"), "completed")
|
||
kinds = ["dispatch" if "dispatch" in c else "run"
|
||
for c in self.calls]
|
||
self.assertEqual(kinds, ["dispatch", "run"])
|
||
run_cmd = [c for c in self.calls if "run" in c][0]
|
||
self.assertEqual(run_cmd[run_cmd.index("--task") + 1], "A")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|