"""云端 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: `，与之前 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 已存在: ，先换名或 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()