cloud: worker dispatch, DB abstraction, packaging 0.1.0, repo CI workflow

cloud/tests/test_worker.py

@@ -0,0 +1,295 @@
+"""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_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()