host
fb3a311da8
此前订阅与分发框架都在，但分发端点传 sender=None， webhook/slack 永远报无 sender、从未真发过。本轮： - push：send_webhook/send_slack（标准库 urllib POST， 单次 10 秒超时，不重试，失败记 errors）；sender=None 时用内置端，显式 sender 优先；webpush/apns/fcm 仍报 缺凭证（原因写明：webpush 需浏览器端三字段订阅与标准库 没有的 P-256 ECDH + AES-GCM；apns 需 HTTP/2；fcm 需 Firebase 服务账号，三者均无现成接收方，不造假发送） - worker：run_one 终态（completed/failed，缺 key 与仓库 失败两条早退路径也在内）后 best-effort 调分发，通知 失败只记日志、绝不影响任务状态；取消（interrupted） 不在事件集合里，不通知 - 测试：真实发送端 6 个（本地回环服务器）、worker 终态 通知 2 个，旧无 sender 断言按新语义改写
393 lines
16 KiB
Python
393 lines
16 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):
|
||
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=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)
|
||
|
||
# -- 取消 ------------------------------------------------------------
|
||
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")
|
||
|
||
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。
|
||
self.assertEqual(row[0], "gpt-6.1-sol")
|
||
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"])
|
||
worker = self._worker(runner=self._recording_runner(fake),
|
||
router=router)
|
||
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()
|