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

cloud/worker.py

@@ -0,0 +1,492 @@
+"""云端 worker：把控制面的 queued 任务接通到 host.py 执行。
+
+零第三方依赖，只用 Python 标准库；不改 host.py 行为、不动
+~/.codex 配置、不碰 wechat-bridge。
+
+流程（与 server.py 的任务状态机对齐）：
+    queued -> running -> completed / failed / interrupted
+
+1. 认领：pick_next() 用 queue.ready_tasks() 按依赖拓扑 + 优先级
+   排序，跳过依赖未完成（未 completed/merged）的任务；再过
+   queue.budget_exceeded() 预算门，超限不派发、任务保持 queued。
+   claim() 用 `WHERE status='queued'` 的 CAS 更新把任务标 running
+   并记 last_heartbeat / updated_at，避免两个 worker 重复认领。
+2. 缺 key 门（执行前必查）：只检查配置的 env_key 名在进程环境
+   里是否存在且非空，绝不读出、打印、写进文件或日志 key 的值。
+   缺 key 时任务直接标 failed，错误固定为
+   `Missing environment variable: <env_key 名>`，与之前
+   real-nightly-tests automation 的失败模式一致，绝不静默假装跑了。
+   env_key 名的来源：显式参数 > codex 配置文件里的 env_key 行
+   （只读名字）> 默认 XCM_API_KEY（中转站 relay 的约定名）。
+3. 执行：经 LocalExecutor 调 host.py 真实子命令。dispatch
+   （建 host 任务并跑一轮 turn）是首次派发；dispatch 报
+   “任务已存在”时改用 run（在已有 host 任务上再跑一轮）续跑。
+   host 任务名由 cloud 任务 id 派生（见 sandbox._host_name），
+   同一任务每次算出同一名字，续跑才能找到同一个 host 任务。
+4. 写回：server 的 tasks 表没有 result/error 列，所以执行结果
+   的持久化分三处，各复用已有结构、不另起字段名：
+   - tasks 表：最终状态 + updated_at（取消时为 interrupted）；
+   - cost.py CostStore：每轮实际执行记一条 source="estimated"
+     的成本记录（host 协议不暴露真实 token 数，估算绝不冒充
+     actual；total_tokens 由 CostStore.record 按 input+output 算）；
+   - models.py：可 fallback 的失败（如 host 状态 timeout）且
+     router 给出下一个模型时，用 record_fallback() 在任务记录上
+     记 model_fallback 事件，并把 fallback_count 写进成本记录。
+     注意：worker 只记录 fallback 事件，不自动用下一个模型重试——
+     自动重试等于未经批准烧 relay 额度，重试与否由调用方决定。
+   审计：认领/完成/失败/取消/缺 key 各记一条 AuthStore 审计
+   （user_id="system"，与 server.sweep_stale_tasks 的做法一致），
+   审计失败不影响任务状态（best-effort）。
+
+取消：host.py 的 argparse 子命令表里没有 interrupt（turn/interrupt
+只存在于 host.py Server 类内部的 app-server 协议调用，未暴露为
+CLI），所以 worker 的 cancel() 走 server 的 cancel 语义：把
+queued/running 任务直接标 interrupted（已终态的任务不动）。
+进行中的 host.py 子进程不会因此被杀（中途杀进程可能把 host
+state 写一半）；run_one() 在执行返回后会先查任务是否已被取消，
+已被取消的一律保持 interrupted，不用执行结果覆盖它。
+
+可注入：runner（假 CLI，见 sandbox.LocalExecutor）与 env 映射
+均可注入，单测不跑真实模型回合、不联网、不登录任何服务。
+"""
+
+import argparse
+import json
+import os
+import re
+import time
+
+from cloud.auth import AuthStore
+from cloud.cost import CostStore, estimate_tokens
+from cloud.db import connect as _db_connect
+from cloud.models import ModelRouter, is_fallbackable, record_fallback
+from cloud.queue import budget_exceeded, ready_tasks
+from cloud.sandbox import LocalExecutor
+
+# 与 server.TERMINAL_STATUSES 对齐：这些状态的任务不再被认领或取消。
+TERMINAL_STATUSES = frozenset(
+    {"completed", "failed", "interrupted", "merged", "merge-failed"})
+COMPLETED_STATUSES = frozenset({"completed", "merged"})
+DEFAULT_ENV_KEY = "XCM_API_KEY"
+DEFAULT_MODEL = "gpt-6.1-sol"
+# host.py dispatch 的原话是 "task 已存在: <name>，先换名或 discard 它"
+#（英文 task + 已存在），取共有子串做判断，不依赖标点。
+_ALREADY_EXISTS_MARKER = "已存在"
+
+
+def _missing_key_error(env_key):
+    """缺 key 的标准错误：只含 env_key 名，与 relay 侧的报错一致。"""
+    return "Missing environment variable: %s" % env_key
+
+
+def resolve_env_key(env_key=None, codex_config_path=None):
+    """确定模型 key 对应的环境变量名（只返回名字，从不返回值）。
+
+    显式参数优先；其次读 codex 配置文件（如 relay 的 config.toml）
+    里的 `env_key = "..."` 行，只取名字；文件不存在或无此行时
+    用默认名。绝不读 ~/.codex/config.toml 之外的默认位置——
+    调用方要读哪个配置文件，必须显式传 codex_config_path。
+    """
+    if env_key:
+        return env_key
+    if codex_config_path:
+        try:
+            with open(codex_config_path, encoding="utf-8") as f:
+                text = f.read()
+        except (OSError, UnicodeDecodeError):
+            text = ""
+        match = re.search(
+            r'^\s*env_key\s*=\s*"([^"]+)"', text, re.MULTILINE)
+        if match:
+            return match.group(1)
+    return DEFAULT_ENV_KEY
+
+
+def _resolve_model(model=None, config_path=None):
+    """模型名：显式参数 > host 配置文件里的 model > host 默认值。"""
+    if model:
+        return model
+    if config_path:
+        try:
+            with open(config_path, encoding="utf-8") as f:
+                cfg = json.load(f)
+        except (OSError, ValueError):
+            cfg = {}
+        if isinstance(cfg, dict) and cfg.get("model"):
+            return cfg["model"]
+    return DEFAULT_MODEL
+
+
+def _task_from_row(row):
+    depends_on = row["depends_on"]
+    try:
+        depends_on = json.loads(depends_on) if depends_on else []
+    except (TypeError, ValueError):
+        depends_on = []
+    if not isinstance(depends_on, list):
+        depends_on = [depends_on]
+    return {
+        "id": row["id"],
+        "project_id": row["project_id"],
+        "prompt": row["prompt"],
+        "status": row["status"],
+        "idempotency_key": row["idempotency_key"],
+        "depends_on": depends_on,
+        "priority": row["priority"],
+        "created_by": row["created_by"],
+        "created_at": row["created_at"],
+        "updated_at": row["updated_at"],
+        "last_heartbeat": row["last_heartbeat"],
+    }
+
+
+class Worker:
+    """单机 worker：一轮认领一个 queued 任务并执行到底。
+
+    前置条件：db_path 须已由 cloud.server.init_db() 建好 tasks 表；
+    本模块不自建 server 的表（成本表与审计表由 CostStore/AuthStore
+    在首次使用时自建，与 server 的做法一致）。
+    """
+
+    def __init__(self, db_path, host_py_path, config_path=None, repo=None,
+                 runner=None, env=None, env_key=None,
+                 codex_config_path=None, model=None, router=None,
+                 max_turns=None, max_cost_usd=None, executor=None,
+                 cost_store=None):
+        self.db_path = str(db_path)
+        self.host_py_path = host_py_path
+        self.config_path = config_path
+        self.repo = repo
+        # env 只用于 key 存在性检查：默认接进程环境，单测可注入
+        # 一个普通 dict；worker 永远不把 env 里的值写出。
+        self._env = env if env is not None else os.environ
+        self.env_key = resolve_env_key(env_key, codex_config_path)
+        self.model = _resolve_model(model, config_path)
+        if router is not None and not isinstance(router, ModelRouter):
+            raise ValueError("router 必须是 cloud.models.ModelRouter")
+        self.router = router
+        self.max_turns = max_turns
+        self.max_cost_usd = max_cost_usd
+        self.usage_turns = 0
+        self.usage_cost_usd = 0.0
+        if executor is not None:
+            self.executor = executor
+        else:
+            self.executor = LocalExecutor(
+                host_py_path, config_path=config_path, runner=runner,
+                repo=repo)
+        self._cost_store = cost_store
+        self._cancel_requested = set()
+
+    # -- 成本存储（惰性建表，避免只认领不执行时也建表）--------------
+    @property
+    def costs(self):
+        if self._cost_store is None:
+            self._cost_store = CostStore(self.db_path)
+        return self._cost_store
+
+    def close(self):
+        """关闭惰性建的成本存储连接（注入的 cost_store 由调用方管）。"""
+        if self._cost_store is not None:
+            try:
+                self._cost_store._conn.close()
+            except Exception:
+                pass
+
+    # -- DB 读写 ------------------------------------------------------
+    def load_tasks(self):
+        conn = _db_connect(self.db_path)
+        try:
+            rows = conn.execute(
+                "SELECT * FROM tasks ORDER BY created_at, id").fetchall()
+        finally:
+            conn.close()
+        return [_task_from_row(r) for r in rows]
+
+    def get_task(self, task_id):
+        conn = _db_connect(self.db_path)
+        try:
+            row = conn.execute(
+                "SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
+        finally:
+            conn.close()
+        return _task_from_row(row) if row is not None else None
+
+    def _claim(self, task_id):
+        """CAS 认领：仅当任务仍为 queued 时标 running 并记心跳。"""
+        now = time.time()
+        conn = _db_connect(self.db_path)
+        try:
+            cur = conn.execute(
+                "UPDATE tasks SET status = 'running',"
+                " last_heartbeat = ?, updated_at = ?"
+                " WHERE id = ? AND status = 'queued'",
+                (now, now, task_id),
+            )
+            conn.commit()
+            return cur.rowcount == 1
+        finally:
+            conn.close()
+
+    def _set_status(self, task_id, status):
+        now = time.time()
+        conn = _db_connect(self.db_path)
+        try:
+            conn.execute(
+                "UPDATE tasks SET status = ?, last_heartbeat = ?,"
+                " updated_at = ? WHERE id = ?",
+                (status, now, now, task_id),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    def heartbeat(self, task_id):
+        """刷新 running 任务的心跳；任务不在 running 时返回 False。"""
+        now = time.time()
+        conn = _db_connect(self.db_path)
+        try:
+            cur = conn.execute(
+                "UPDATE tasks SET last_heartbeat = ?, updated_at = ?"
+                " WHERE id = ? AND status = 'running'",
+                (now, now, task_id),
+            )
+            conn.commit()
+            return cur.rowcount == 1
+        finally:
+            conn.close()
+
+    def _audit(self, project_id, action, detail=""):
+        """best-effort 审计：失败不影响任务状态。"""
+        try:
+            store = AuthStore(self.db_path)
+            try:
+                store.audit("system", project_id, action, detail)
+            finally:
+                try:
+                    store.conn.close()
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+    # -- 认领 ----------------------------------------------------------
+    def _budget_exceeded(self):
+        return budget_exceeded(self.usage_turns, self.max_turns,
+                               self.usage_cost_usd, self.max_cost_usd)
+
+    def pick_next(self):
+        """下一个可派任务，或 None（无就绪任务 / 预算超限 / 依赖成环）。
+
+        依赖满足的定义：依赖任务状态为 completed 或 merged。
+        依赖成环时 ready_tasks 抛 ValueError，worker 视为无就绪任务
+        （成环任务本就永远派不出去，run_all 不应因此崩溃）。
+        """
+        if self._budget_exceeded():
+            return None
+        tasks = self.load_tasks()
+        completed_ids = {t["id"] for t in tasks
+                         if t["status"] in COMPLETED_STATUSES}
+        try:
+            ready = ready_tasks(tasks, completed_ids)
+        except ValueError:
+            return None
+        return ready[0] if ready else None
+
+    # -- 取消（server cancel 语义，见模块 docstring）-------------------
+    def cancel(self, task_id):
+        """把 queued/running 任务标 interrupted；已终态或不存在返回 False。"""
+        task = self.get_task(task_id)
+        if task is None or task["status"] in TERMINAL_STATUSES:
+            return False
+        self._cancel_requested.add(task_id)
+        conn = _db_connect(self.db_path)
+        try:
+            cur = conn.execute(
+                "UPDATE tasks SET status = 'interrupted', updated_at = ?"
+                " WHERE id = ? AND status NOT IN"
+                " ('completed', 'failed', 'interrupted', 'merged',"
+                " 'merge-failed')",
+                (time.time(), task_id),
+            )
+            conn.commit()
+            claimed = cur.rowcount == 1
+        finally:
+            conn.close()
+        if claimed:
+            self._audit(task["project_id"], "worker_cancelled", task_id)
+        return claimed
+
+    def _was_cancelled(self, task_id):
+        if task_id in self._cancel_requested:
+            return True
+        task = self.get_task(task_id)
+        return task is not None and task["status"] == "interrupted"
+
+    # -- 缺 key 门 ------------------------------------------------------
+    def key_available(self):
+        """只查存在性：名在环境里且值非空即 True；值本身绝不返回。"""
+        try:
+            value = self._env.get(self.env_key)
+        except Exception:
+            return False
+        return value is not None and value != ""
+
+    # -- 执行 ----------------------------------------------------------
+    def run_one(self, task=None):
+        """认领并执行一个任务，返回结果 dict（见模块 docstring）。
+
+        未派发时返回 {"dispatched": False, "reason": ...} 且不动任务；
+        派发后（即使缺 key 失败）返回 dispatched=True 与最终状态。
+        """
+        if self._budget_exceeded():
+            return {"dispatched": False, "reason": "budget_exceeded"}
+        if task is None:
+            task = self.pick_next()
+        if task is None:
+            return {"dispatched": False, "reason": "no_ready_task"}
+        task_id = task["id"]
+        if not self._claim(task_id):
+            return {"dispatched": False, "reason": "already_claimed",
+                    "task_id": task_id}
+        task = self.get_task(task_id) or task
+        self._audit(task["project_id"], "worker_claim", task_id)
+
+        # 缺 key 门：认领之后、调 host.py 之前；缺 key 直接 failed，
+        # 不调 runner（假 CLI 单测可断言 runner 未被调用）。
+        if not self.key_available():
+            error = _missing_key_error(self.env_key)
+            self._set_status(task_id, "failed")
+            self._audit(task["project_id"], "worker_missing_key", error)
+            return {"dispatched": True, "task_id": task_id,
+                    "host_name": self.executor._host_name(task)
+                    if hasattr(self.executor, "_host_name") else task_id,
+                    "status": "failed", "executor": "local",
+                    "stdout": "", "error": error, "events": [],
+                    "cost_record_id": None}
+
+        events = []
+        fallback_count = 0
+        result = None
+        error = None
+        try:
+            result = self.executor.dispatch_task(task)
+            combined = (result.get("stdout", "") or "") + \
+                (result.get("stderr", "") or "")
+            if result["status"] != "completed" and \
+                    _ALREADY_EXISTS_MARKER in combined:
+                # host 任务已存在（上轮 worker 在 dispatch 后、写回前
+                # 崩溃）：改用 run 在已有任务上续跑，不重复建 worktree。
+                result = self.executor.run_task(task)
+        except Exception as exc:  # runner 异常也如实写回 failed
+            result = {"executor": "local", "task_id": task_id,
+                      "status": "failed", "stdout": "", "returncode": None,
+                      "host_status": None, "stderr": ""}
+            error = str(exc)
+
+        # 取消优先：执行期间被 cancel 的任务保持 interrupted，
+        # 绝不用执行结果覆盖它（见模块 docstring 的取消说明）。
+        if self._was_cancelled(task_id):
+            self._set_status(task_id, "interrupted")
+            return {"dispatched": True, "task_id": task_id,
+                    "host_name": self.executor._host_name(task)
+                    if hasattr(self.executor, "_host_name") else task_id,
+                    "status": "interrupted", "executor": "local",
+                    "stdout": (result or {}).get("stdout", ""),
+                    "error": "cancelled", "events": [],
+                    "cost_record_id": None}
+
+        stdout = result.get("stdout", "") or ""
+        host_status = result.get("host_status")
+        final = result["status"]
+        if final != "completed":
+            if error is None:
+                error = stdout.strip() or "host execution failed"
+            # fallback 事件：只记、不自动重试（见模块 docstring）。
+            if self.router is not None and host_status is not None \
+                    and is_fallbackable(host_status):
+                nxt = self.router.next_after(self.model, host_status)
+                if nxt is not None:
+                    record = dict(task)
+                    record["events"] = events
+                    # record_fallback 本身会 append 到 events，
+                    # 这里不要再 append 一次，否则事件会重复。
+                    record_fallback(
+                        record, self.model, nxt, str(host_status))
+                    fallback_count = 1
+                    self._audit(task["project_id"], "worker_fallback",
+                                "%s -> %s" % (self.model, nxt))
+        self._set_status(task_id, final)
+        self.usage_turns += 1
+        # 成本：host 协议不暴露真实 token 数，用 cost.py 的估算结构
+        # 记 estimated，绝不标 actual；fallback_count 复用同一记录。
+        total = estimate_tokens(task.get("prompt") or "", stdout)
+        cost_id = self.costs.record(
+            task_id, task["project_id"], self.model,
+            input_tokens=total, output_tokens=0, usd_micros=None,
+            source="estimated", fallback_count=fallback_count)
+        self._audit(task["project_id"],
+                    "worker_completed" if final == "completed"
+                    else "worker_failed", task_id)
+        return {"dispatched": True, "task_id": task_id,
+                "host_name": self.executor._host_name(task)
+                if hasattr(self.executor, "_host_name") else task_id,
+                "status": final, "executor": "local",
+                "stdout": stdout, "error": error, "events": events,
+                "cost_record_id": cost_id}
+
+    def run_all(self, max_tasks=None):
+        """循环认领执行，直到无就绪任务、预算超限或达到 max_tasks。
+
+        每轮都重新读库，所以上一轮刚 completed 的任务会解锁它的
+        依赖任务；每轮至少推进一个任务到终态，不会空转死循环。
+        """
+        results = []
+        while max_tasks is None or len(results) < max_tasks:
+            result = self.run_one()
+            if not result["dispatched"]:
+                break
+            results.append(result)
+        return results
+
+
+def main(argv=None):
+    """单次派发循环：python3 -m cloud.worker --db <库> --host-py <路径>。
+
+    只打印任务结果摘要（状态、host 名、错误文本），绝不打印
+    任何环境变量的值。
+    """
+    parser = argparse.ArgumentParser(description="cloud worker (one pass)")
+    parser.add_argument("--db", default="cloud-control-plane.db")
+    parser.add_argument("--host-py", default=None,
+                        help="host.py 路径，缺省为 cloud/ 上一级的 host.py")
+    parser.add_argument("--config", default=None)
+    parser.add_argument("--repo", default=None)
+    parser.add_argument("--env-key", default=None)
+    parser.add_argument("--codex-config", default=None,
+                        help="含 env_key 行的 codex 配置文件（只读名字）")
+    parser.add_argument("--model", default=None)
+    parser.add_argument("--max-turns", type=int, default=None)
+    parser.add_argument("--max-tasks", type=int, default=None)
+    args = parser.parse_args(argv)
+    host_py = args.host_py
+    if host_py is None:
+        here = os.path.dirname(os.path.abspath(__file__))
+        host_py = os.path.join(os.path.dirname(here), "host.py")
+    worker = Worker(args.db, host_py, config_path=args.config,
+                    repo=args.repo, env_key=args.env_key,
+                    codex_config_path=args.codex_config, model=args.model,
+                    max_turns=args.max_turns)
+    try:
+        results = worker.run_all(max_tasks=args.max_tasks)
+    finally:
+        worker.close()
+    print(json.dumps(
+        [{"task_id": r.get("task_id"), "status": r.get("status"),
+          "reason": r.get("reason"), "error": r.get("error")}
+         for r in results],
+        indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()