codex-app-server/cloud/worker.py

"""云端 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）。

取消：worker 的 cancel() 走 server 的 cancel 语义：把
queued/running 任务直接标 interrupted（已终态的任务不动），
且正在跑的子进程会真杀——子进程以进程组启动（start_new_session），
取消时先 SIGTERM 给整个进程组、等 2 秒、仍活着再 SIGKILL，
杀完后任务保持 interrupted，执行结果不覆盖它。run_one() 在
执行返回后会先查任务是否已被取消，已被取消的一律保持
interrupted，不用执行结果覆盖它。

可注入：runner（假 CLI，见 sandbox.LocalExecutor）与 env 映射
均可注入，单测不跑真实模型回合、不联网、不登录任何服务。

执行器选择：Worker 默认用 LocalExecutor；传入 DockerExecutor
（或 CLI --executor docker）时，派发走容器内 host.py，取消时
按任务派生的容器名 docker kill，不只改数据库状态；本机无
docker 时 DockerExecutor 在调 runner 前就明确报错，任务标
failed 且 error_text 含 docker 缺失，绝不静默降级为本地执行。
"""

import argparse
import json
import os
import re
import signal
import subprocess
import sys
import time
from types import SimpleNamespace

from cloud.auth import AuthStore
from cloud.cost import CostStore, estimate_tokens
from cloud.db import connect as _db_connect
from cloud.logging_config import get_logger, log_event, setup_logging
from cloud.models import ModelRouter, is_fallbackable, record_fallback
from cloud.queue import budget_exceeded, ready_tasks
from cloud.sandbox import DockerExecutor, LocalExecutor
from cloud.sandbox import _default_runner as _sandbox_default_runner

# 与 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, process_factory=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()
        # 真取消用：正在跑的子进程（Popen 对象）与任务 id。
        # process_factory 可注入：单测用它返回一个假长跑子进程，
        # 真实执行时走 LocalExecutor.popen_task（进程组启动）。
        self.process_factory = process_factory
        self._running_proc = None
        self._running_task_id = None
        self._running_container = None
        self._loop_stop = False
        self.logger = get_logger("cloud.worker")

    # -- 成本存储（惰性建表，避免只认领不执行时也建表）--------------
    @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 _ensure_result_columns(self, conn):
        """worker 侧的迁移兜底：与 server.migrate_tasks_table 同列。

        worker 的库理论上已由 server.init_db 迁过；这里再查一次
        是为了 worker 直连老库（测试里手工建的表）时也不炸。
        """
        try:
            cols = {row[1] for row in conn.execute(
                "PRAGMA table_info(tasks)").fetchall()}
        except Exception:
            return
        for name, coltype in {
                "result_text": "TEXT", "error_text": "TEXT",
                "actual_model": "TEXT", "started_at": "REAL",
                "finished_at": "REAL", "child_pid": "INTEGER"}.items():
            if name not in cols:
                try:
                    conn.execute(
                        "ALTER TABLE tasks ADD COLUMN %s %s"
                        % (name, coltype))
                except Exception:
                    pass

    def _write_execution_result(self, task_id, result_text=None,
                                error_text=None, actual_model=None,
                                started_at=None, finished_at=None,
                                child_pid=None):
        """把执行结果写进 tasks 新列；老库缺列时先迁移再写。"""
        conn = _db_connect(self.db_path)
        try:
            self._ensure_result_columns(conn)
            if child_pid is None:
                conn.execute(
                    "UPDATE tasks SET result_text = ?, error_text = ?,"
                    " actual_model = ?, started_at = ?, finished_at = ?,"
                    " updated_at = ? WHERE id = ?",
                    ((result_text or "")[:4000] if result_text else None,
                     (error_text or "")[:4000] if error_text else None,
                     actual_model, started_at, finished_at,
                     time.time(), task_id),
                )
            else:
                conn.execute(
                    "UPDATE tasks SET result_text = ?, error_text = ?,"
                    " actual_model = ?, started_at = ?, finished_at = ?,"
                    " child_pid = ?, updated_at = ? WHERE id = ?",
                    ((result_text or "")[:4000] if result_text else None,
                     (error_text or "")[:4000] if error_text else None,
                     actual_model, started_at, finished_at, child_pid,
                     time.time(), task_id),
                )
            conn.commit()
        finally:
            conn.close()

    def _set_child_pid(self, task_id, pid):
        conn = _db_connect(self.db_path)
        try:
            self._ensure_result_columns(conn)
            conn.execute(
                "UPDATE tasks SET child_pid = ? WHERE id = ?",
                (pid, task_id),
            )
            conn.commit()
        finally:
            conn.close()

    def _kill_proc_group(self, proc):
        """真杀：先 SIGTERM 给整个进程组，等 2 秒，仍活着再 SIGKILL。

        proc 须是 start_new_session 启动的 Popen（自成进程组），
        杀组才能连同它拉起的子孙一起杀掉，不只杀父进程。
        """
        try:
            pgid = os.getpgid(proc.pid)
        except Exception:
            return
        try:
            os.killpg(pgid, signal.SIGTERM)
        except (ProcessLookupError, PermissionError, OSError):
            return
        deadline = time.time() + 2.0
        while time.time() < deadline:
            try:
                os.killpg(pgid, 0)
            except (ProcessLookupError, PermissionError, OSError):
                return
            time.sleep(0.05)
        try:
            os.killpg(pgid, signal.SIGKILL)
        except (ProcessLookupError, PermissionError, OSError):
            pass

    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。

        正在跑的子进程要真杀：先杀进程组（见 _kill_proc_group），
        再标状态；已终态或不存在返回 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)
        if self._running_task_id == task_id and self._running_proc is not None:
            self._kill_proc_group(self._running_proc)
            try:
                self._running_proc.wait(timeout=3)
            except Exception:
                pass
            self._running_proc = None
            self._running_task_id = None
        # Docker 取消：按任务派生的容器名真杀，不只改数据库状态。
        # 同步路径（runner 阻塞在 docker run 里）下 _running_proc
        # 为空，此时靠 _running_container 杀容器；kill 尽力而为，
        # 失败也不影响后续的状态标记。
        if self._running_task_id == task_id and \
                self._running_container is not None and \
                hasattr(self.executor, "kill_container"):
            try:
                self.executor.kill_container(self._running_container)
            except Exception:
                pass
            self._running_container = None
        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 not claimed:
            # 并发竞态：run_one 的轮询已看到本次取消的标记、先一步
            # 把任务标成 interrupted——中断确实由本次取消触发，返回 True。
            current = self.get_task(task_id)
            if current is not None and current["status"] == "interrupted":
                claimed = True
        if claimed:
            self._audit(task["project_id"], "worker_cancelled", task_id)
            log_event(self.logger, "worker_cancelled", task_id=task_id,
                      project_id=task["project_id"], status="interrupted")
        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 _executor_name(self):
        """执行器名：DockerExecutor -> docker，其它 -> local。"""
        if isinstance(self.executor, DockerExecutor):
            return "docker"
        return "local"
    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 _use_popen(self):
        """是否走进程模式：注入了 process_factory，或执行器是
        默认 runner 的 LocalExecutor（真实 subprocess）。"""
        if self.process_factory is not None:
            return True
        return (isinstance(self.executor, LocalExecutor)
                and getattr(self.executor, "runner", None)
                is _sandbox_default_runner
                and hasattr(self.executor, "popen_task"))

    def _execute_via_process(self, task, kind="dispatch"):
        """进程模式执行一轮：启动、轮询、取消时杀进程组。

        返回与同步路径同款 result dict（status/stdout/stderr/
        host_status/returncode）。取消时返回 status='cancelled'
        的标记，调用方据此保持 interrupted。
        """
        task_id = task["id"]
        if self.process_factory is not None:
            proc = self.process_factory(task)
        else:
            proc = self.executor.popen_task(task, kind=kind)
        self._running_proc = proc
        self._running_task_id = task_id
        try:
            self._set_child_pid(task_id, getattr(proc, "pid", None))
        except Exception:
            pass
        cancelled = False
        try:
            while proc.poll() is None:
                if task_id in self._cancel_requested or self._loop_stop:
                    cancelled = True
                    self._kill_proc_group(proc)
                    break
                time.sleep(0.05)
            try:
                out, err = proc.communicate(timeout=5)
            except Exception:
                try:
                    proc.kill()
                except Exception:
                    pass
                out, err = proc.communicate()
        finally:
            self._running_proc = None
            self._running_task_id = None
        if cancelled:
            return {"executor": "local", "task_id": task_id,
                    "status": "cancelled", "stdout": out or "",
                    "stderr": err or "", "returncode": proc.returncode,
                    "host_status": None}
        fake = SimpleNamespace(returncode=proc.returncode,
                               stdout=out or "", stderr=err or "")
        if hasattr(self.executor, "_result_from_runner"):
            return self.executor._result_from_runner(task_id, fake)
        return {"executor": "local", "task_id": task_id,
                "status": "completed" if proc.returncode == 0 else "failed",
                "stdout": out or "", "stderr": err or "",
                "returncode": proc.returncode, "host_status": None}

    # -- 执行 ----------------------------------------------------------
    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)
        log_event(self.logger, "worker_claim", task_id=task_id,
                  project_id=task["project_id"], status="running")
        started_at = time.time()

        # 缺 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._write_execution_result(
                task_id, error_text=error, actual_model=self.model,
                started_at=started_at, finished_at=time.time())
            self._audit(task["project_id"], "worker_missing_key", error)
            log_event(self.logger, "worker_missing_key", task_id=task_id,
                      project_id=task["project_id"], status="failed")
            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": self._executor_name(),
                    "stdout": "", "error": error, "events": [],
                    "cost_record_id": None}

        events = []
        fallback_count = 0
        result = None
        error = None
        # Docker 同步路径：dispatch 阻塞在 docker run 里时，取消
        # 线程靠 _running_container 按名杀容器（见 cancel()）。
        if hasattr(self.executor, "container_name"):
            try:
                self._running_container = \
                    self.executor.container_name(task)
                self._running_task_id = task_id
            except Exception:
                self._running_container = None
        try:
            if self._use_popen():
                result = self._execute_via_process(task)
                combined = (result.get("stdout", "") or "") + \
                    (result.get("stderr", "") or "")
                if result["status"] not in ("completed", "cancelled") and \
                        _ALREADY_EXISTS_MARKER in combined:
                    result = self._execute_via_process(task, kind="run")
            else:
                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:
                    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)
        # 同步路径执行结束：清掉取消用的容器标记（取消路径里
        # 若已清过，这里再清一次无害）。
        self._running_container = None
        if self._running_task_id == task_id and \
                self._running_proc is None:
            self._running_task_id = None

        # 取消优先：执行期间被 cancel 的任务保持 interrupted，
        # 绝不用执行结果覆盖它（见模块 docstring 的取消说明）。
        if self._was_cancelled(task_id) or result.get("status") == "cancelled":
            self._set_status(task_id, "interrupted")
            self._write_execution_result(
                task_id, result_text=(result or {}).get("stdout", ""),
                error_text="cancelled", actual_model=self.model,
                started_at=started_at, finished_at=time.time())
            log_event(self.logger, "worker_cancelled", task_id=task_id,
                      project_id=task["project_id"], status="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": self._executor_name(),
                    "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)
        finished_at = time.time()
        self._write_execution_result(
            task_id, result_text=stdout, error_text=error,
            actual_model=self.model, started_at=started_at,
            finished_at=finished_at)
        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)
        log_event(self.logger,
                  "worker_completed" if final == "completed"
                  else "worker_failed",
                  task_id=task_id, project_id=task["project_id"],
                  status=final)
        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": self._executor_name(),
                "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 run_loop(self, interval=300, max_seconds=None, stop_event=None):
        """常驻循环：每 interval 秒跑一轮 run_all，空闲也不退出。

        退出条件：stop_event 被置位、_loop_stop 被信号处理置位、
        或 max_seconds 到期。退出时若有正在跑的任务（进程模式下
        _running_task_id 非空），先杀进程组、再标 interrupted，
        不让任务永远卡在 running。
        """
        started = time.time()
        all_results = []
        try:
            while not self._loop_stop:
                if stop_event is not None and stop_event.is_set():
                    break
                if max_seconds is not None and \
                        time.time() - started >= max_seconds:
                    break
                try:
                    results = self.run_all()
                except Exception as exc:
                    # 常驻循环一轮炸了不带走整个 worker：记结构化日志，
                    # 睡完本轮间隔再试下一轮（缺表、锁冲突都属此类）。
                    log_event(self.logger, "worker_loop_error",
                              message="worker_loop_error: %s" % exc)
                    results = []
                all_results.extend(results)
                # 按 0.2 秒切片睡，期间可响应 stop_event/_loop_stop。
                deadline = time.time() + interval
                while time.time() < deadline:
                    if self._loop_stop:
                        break
                    if stop_event is not None and stop_event.is_set():
                        break
                    if max_seconds is not None and \
                            time.time() - started >= max_seconds:
                        break
                    time.sleep(0.2)
        finally:
            if self._running_task_id is not None:
                task_id = self._running_task_id
                if self._running_proc is not None:
                    self._kill_proc_group(self._running_proc)
                    try:
                        self._running_proc.wait(timeout=3)
                    except Exception:
                        pass
                self._running_proc = None
                self._running_task_id = None
                self._set_status(task_id, "interrupted")
        return all_results


def _build_cli_executor(args, host_py):
    """CLI 的执行器构造：local 时返回 None（Worker 默认 LocalExecutor），
    docker 时返回 DockerExecutor（env_key 名用 Worker 的解析规则，
    值本身绝不从这里传——DockerExecutor 只查存在性，-e 只传名）。"""
    if getattr(args, "executor", "local") != "docker":
        return None
    return DockerExecutor(
        host_py, config_path=args.config, image=args.docker_image,
        cpus=args.docker_cpus, memory=args.docker_memory,
        user=args.docker_user, worktrees_root=args.docker_worktrees_root,
        env_key=resolve_env_key(args.env_key, args.codex_config),
        timeout=args.docker_timeout,
        extra_mounts=args.docker_extra_mounts)


def main(argv=None):
    """worker CLI：默认单次 run_all；--loop 进常驻循环。

    只打印任务结果摘要（状态、host 名、错误文本），绝不打印
    任何环境变量的值。
    """
    parser = argparse.ArgumentParser(description="cloud worker")
    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)
    parser.add_argument("--executor", default="local",
                        choices=["local", "docker"],
                        help="执行器：local 调 host.py 子进程；docker"
                             " 调 DockerExecutor（无 docker 时任务"
                             "明确失败，不静默降级）")
    parser.add_argument("--docker-image", default="codex-app-server:latest")
    parser.add_argument("--docker-cpus", default="2")
    parser.add_argument("--docker-memory", default="2g")
    parser.add_argument("--docker-user", default="1000:1000")
    parser.add_argument("--docker-worktrees-root", default=None)
    parser.add_argument("--docker-timeout", type=float, default=None)
    parser.add_argument("--docker-extra-mounts", default=None,
                        help="docker 额外挂载，逗号分隔 host:container[:ro|rw]，"
                             "用于容器内 host.py 写 work/、host-state.json，"
                             "以及挂载原仓库和 codex-home（同路径挂载，"
                             "容器内配置无需改路径）")
    parser.add_argument("--loop", action="store_true",
                        help="常驻循环，每 --interval 秒跑一轮")
    parser.add_argument("--interval", type=int, default=300)
    parser.add_argument("--max-seconds", type=float, default=None)
    args = parser.parse_args(argv)
    setup_logging()
    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,
                    executor=_build_cli_executor(args, host_py))
    # CLI 起 worker 前先建表/迁移：全新库直接起循环不应因缺表而炸。
    try:
        from cloud.server import init_db as _init_db
        _init_db(args.db)
    except Exception as exc:
        print("worker db 初始化失败：%s" % exc, file=sys.stderr)
        raise SystemExit(1)
    if args.loop:
        def _on_signal(signum, frame):
            worker._loop_stop = True
        signal.signal(signal.SIGTERM, _on_signal)
        signal.signal(signal.SIGINT, _on_signal)
    try:
        if args.loop:
            results = worker.run_loop(interval=args.interval,
                                      max_seconds=args.max_seconds)
        else:
            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()