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，不用执行结果覆盖它。

任务级超时与失败重试（生产级，见 Worker.__init__ 参数）：
- task_timeout（默认 1800 秒，CLI --task-timeout 可调）是
  worker 级的每任务截止时间，与下面两个参数不要混淆：
  --docker-timeout 是执行器级的单次 docker run 超时（Docker
  Executor 自身的 timeout），--max-seconds 是 run_loop 的
  整体墙钟上限（不是每任务）。三者互不覆盖：Docker 路径下
  worker 取 min(task_timeout, docker_timeout)（docker_timeout
  未设时取 task_timeout）作为本次 docker run 的超时，超时后
  真杀容器、任务置 failed（超时是失败，不是取消，绝不标
  interrupted）；Local popen 路径下 worker 在轮询循环里到
  截止时间就杀进程组、同样置 failed；Local 同步路径（假
  runner / 无超时的 subprocess.run）无法在中途抢占一个已阻塞
  的调用，超时只能靠执行器自身的 timeout，这是诚实边界，
  不要假装同步路径也能中途杀掉。
- max_retries（默认 1 次重试，CLI --max-retries 可调，0 表示
  不重试）是对模型任务失败的有限重试：取消不重试；已记
  fallback 事件的失败不做同模型重试（fallback 本身就是模型
  级的重试机制，再同模型重试会重复烧额度）；其它失败（含
  worker 超时）重试前按 retry_backoff_seconds（默认 5 秒，
  测试可注入 0）等待，重试次数记在审计（worker_retry）与
  最终 error_text（[retries=N] 前缀）里，不另起新表；重试
  耗尽后任务才最终 failed，成本记录与 usage_turns 只记一次
  （按最终结果记，不按尝试次数重复记）。
- 容器清理兜底：Worker.cleanup_stale_containers() 列出
  codex- 前缀的残留容器并逐个 kill+rm，失败只记日志不炸
  worker；run_loop 启动时与 CLI 启动时各调一次（run_all /
  __init__ 里不调，避免单测构造 Worker 时多余 I/O）。

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

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

项目仓库：run_one 派发前查任务所属项目的 repo_url（projects
表），非空时经 cloud.repos.resolve_repo 在 repos_dir（CLI
--repos-dir，默认 ./repos）下 clone/fetch 该项目仓库，设
task["repo"] 再交给执行器；Local 与 Docker 执行器都已读
task["repo"]，同一条路径。repo_url 为空时走原有默认 repo；
解析失败时任务直接 failed，错误文本写清是仓库解析失败。
Docker 下 repos_dir 必须放在 extra-mounts 已挂载的目录下，
否则容器内 repo 路径不存在（见 cloud/README.md）。
推送到远端不在本轮范围：任务在本地检出上执行、本地 worktree
合并仍按 host.py 现有流程，git push 需要仓库写凭证，控制面
不存、不推。
"""

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.repos import RepoError, resolve_repo
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,
                 task_timeout=1800, max_retries=1,
                 retry_backoff_seconds=5.0, repos_dir="./repos",
                 repos_runner=None, heartbeat_file=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")
        # 任务级超时 / 失败重试（见模块 docstring 的生产级说明）。
        self.task_timeout = task_timeout
        self.max_retries = max_retries
        self.retry_backoff_seconds = retry_backoff_seconds
        # 项目仓库检出根目录：CLI --repos-dir 可调（默认 ./repos）；
        # Docker 执行器下它必须放在 extra-mounts 已挂载的目录下，
        # 否则容器内 task["repo"] 路径不存在——见 cloud/README.md。
        # repos_runner 可注入（单测用假 git runner），缺省走真实 git。
        self.repos_dir = repos_dir
        self.repos_runner = repos_runner
        # 运维心跳：路径优先显式参数，其次环境变量，默认当前目录下
        # 的 worker-heartbeat.json；主机路径不在代码里硬编码，主机
        # 部署时用 --heartbeat-file 或环境变量传入（见 deploy/README）。
        # 文件里只有 last_seen 与 pid，绝无任何密钥。
        self.heartbeat_file = heartbeat_file or \
            os.environ.get("CLOUD_WORKER_HEARTBEAT_FILE") or \
            "./worker-heartbeat.json"

    def _write_heartbeat(self):
        """写一次运维心跳：写失败只记日志，绝不炸 worker。

        先写临时文件再 os.replace，保证 server 读到的要么是旧的
        完整文件、要么是新的完整文件，不会读到写一半的 JSON。
        """
        try:
            payload = {"last_seen": time.time(), "pid": os.getpid()}
            tmp = "%s.tmp" % self.heartbeat_file
            parent = os.path.dirname(os.path.abspath(self.heartbeat_file))
            os.makedirs(parent, exist_ok=True)
            with open(tmp, "w", encoding="utf-8") as f:
                f.write(json.dumps(payload))
            os.replace(tmp, self.heartbeat_file)
        except Exception as exc:
            log_event(self.logger, "worker_heartbeat_failed",
                      message="worker heartbeat write failed: %s" % exc)

    # -- 成本存储（惰性建表，避免只认领不执行时也建表）--------------
    @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 cleanup_stale_containers(self):
        """启动清理兜底：清掉上一轮残留的 codex- 容器。

        只在执行器有 cleanup_stale_containers 时做（Docker）；
        列出或删除失败只记日志，绝不抛异常炸掉 worker。
        返回已清理的容器名列表（非 Docker 执行器返回 []）。
        """
        cleanup = getattr(self.executor, "cleanup_stale_containers", None)
        if cleanup is None:
            return []
        try:
            cleaned = cleanup()
        except Exception as exc:
            log_event(self.logger, "worker_cleanup_failed",
                      message="cleanup stale containers failed: %s" % exc)
            return []
        if cleaned:
            log_event(self.logger, "worker_cleanup",
                      message="cleaned stale containers: %s"
                              % ",".join(cleaned))
        return cleaned

    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",
                             deadline_seconds=None):
        """进程模式执行一轮：启动、轮询、取消时杀进程组。

        返回与同步路径同款 result dict（status/stdout/stderr/
        host_status/returncode）。取消时返回 status='cancelled'
        的标记，调用方据此保持 interrupted；deadline_seconds
        到期时杀进程组并返回 status='timeout' 的标记，调用方
        据此置 failed（超时是失败，不是取消）。
        """
        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
        timed_out = False
        started = time.time()
        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
                if deadline_seconds is not None and \
                        time.time() - started >= deadline_seconds:
                    timed_out = 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}
        if timed_out:
            return {"executor": "local", "task_id": task_id,
                    "status": "timeout", "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 _project_repo_url(self, project_id):
        """查任务所属项目的 repo_url；查不到项目、列缺失、表缺失时
        一律返回 None（走原有默认 repo，不报错、不炸掉派发）。
        repo_url 本身只在 projects 表里，绝不从任务 prompt 里猜。"""
        try:
            conn = _db_connect(self.db_path)
            try:
                row = conn.execute(
                    "SELECT repo_url FROM projects WHERE id = ?",
                    (project_id,)).fetchone()
            finally:
                conn.close()
        except Exception:
            return None
        if row is None:
            return None
        try:
            url = row["repo_url"]
        except (KeyError, IndexError, TypeError):
            return None
        return url or None

    def _apply_project_repo(self, task):
        """派发前按项目 repo_url 解析仓库，返回错误文本或 None。

        repo_url 为空：不动 task，返回 None（走原有默认 repo）。
        repo_url 非空：调 cloud.repos.resolve_repo 得到本地检出，
        设 task["repo"] 再交给执行器——LocalExecutor 与
        DockerExecutor 的 dispatch_task 都已读 task["repo"]，同一条
        路径，不另起分支。resolve 失败返回错误文本，调用方据此把
        任务置 failed（仓库解析失败绝不标 completed，也不在此重试，
        避免对不可用的仓库反复 clone 烧时间）。
        """
        repo_url = self._project_repo_url(task.get("project_id"))
        if not repo_url:
            return None
        try:
            task["repo"] = resolve_repo(
                repo_url, self.repos_dir, task.get("project_id"),
                runner=self.repos_runner)
        except (RepoError, ValueError) as exc:
            return "仓库解析失败: %s" % exc
        except Exception as exc:
            return "仓库解析失败: %s" % exc
        return None

    # -- 执行 ----------------------------------------------------------
    def _run_attempt(self, task):
        """单次执行尝试（不含认领、key 门、重试、成本与最终审计）。

        返回 (result, error)：result 是执行器结果 dict，error 是
        runner 异常时的文本（无异常时 None）。Docker 路径下把
        worker 级 task_timeout 与执行器级 docker_timeout 取较小
        值作为本次 docker run 超时（docker_timeout 未设时取
        task_timeout）；Local popen 路径下把 task_timeout 作为
        轮询截止时间；Local 同步路径无法中途抢占，超时只能靠
        执行器自身 timeout（见模块 docstring 的诚实边界）。
        调用方负责：取消优先、超时转 failed、重试、成本与审计。
        """
        task_id = task["id"]
        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
        # Docker 超时取较小值：两个超时都不应静默延长另一个。
        saved_docker_timeout = None
        is_docker = isinstance(self.executor, DockerExecutor)
        if is_docker and self.task_timeout is not None:
            saved_docker_timeout = self.executor.timeout
            if saved_docker_timeout is None:
                self.executor.timeout = self.task_timeout
            else:
                try:
                    self.executor.timeout = min(
                        float(self.task_timeout),
                        float(saved_docker_timeout))
                except (TypeError, ValueError):
                    self.executor.timeout = self.task_timeout
        try:
            if self._use_popen():
                result = self._execute_via_process(
                    task, deadline_seconds=self.task_timeout)
                combined = (result.get("stdout", "") or "") + \
                    (result.get("stderr", "") or "")
                if result["status"] not in (
                        "completed", "cancelled", "timeout") and \
                        _ALREADY_EXISTS_MARKER in combined:
                    result = self._execute_via_process(
                        task, kind="run",
                        deadline_seconds=self.task_timeout)
            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)
        finally:
            if is_docker:
                self.executor.timeout = saved_docker_timeout
            # 同步路径执行结束：清掉取消用的容器标记（取消路径里
            # 若已清过，这里再清一次无害）。
            self._running_container = None
            if self._running_task_id == task_id and \
                    self._running_proc is None:
                self._running_task_id = None
        # Docker 超时标记：执行器已在超时时 kill+rm 容器并返回
        # failed（stdout 含 timeout），这里把状态统一成 timeout，
        # 调用方据此置 failed 并记 worker_timeout 审计。
        if result is not None and result.get("status") == "failed" and \
                "docker run timeout" in (result.get("stdout", "") or ""):
            result = dict(result, status="timeout")
        return result, error

    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}

        # 项目仓库解析：认领 + 缺 key 门之后、真正执行之前。
        # repo_url 为空时不动（默认 repo）；解析失败时任务直接 failed，
        # 错误文本写清是仓库解析失败，绝不标 completed，也不在此重试。
        repo_error = self._apply_project_repo(task)
        if repo_error is not None:
            self._set_status(task_id, "failed")
            self._write_execution_result(
                task_id, error_text=repo_error, actual_model=self.model,
                started_at=started_at, finished_at=time.time())
            self._audit(task["project_id"], "worker_repo_failed", repo_error)
            log_event(self.logger, "worker_repo_failed", 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": repo_error, "events": [],
                    "cost_record_id": None}

        events = []
        fallback_count = 0
        result = None
        error = None
        retry_count = 0
        # 有限重试：取消不重试；已记 fallback 的失败不做同模型重试
        # （fallback 本身就是模型级重试机制）；其它失败（含 worker
        # 超时）按 max_retries 重试，退避等待后重跑一次 _run_attempt。
        for attempt in range(self.max_retries + 1):
            result, error = self._run_attempt(task)
            if self._was_cancelled(task_id) or \
                    result.get("status") == "cancelled":
                break
            if result.get("status") == "completed":
                break
            host_status_try = result.get("host_status")
            if self.router is not None and host_status_try is not None \
                    and is_fallbackable(host_status_try) and \
                    self.router.next_after(
                        self.model, host_status_try) is not None:
                break
            if attempt < self.max_retries:
                retry_count += 1
                self._audit(
                    task["project_id"], "worker_retry",
                    "%s: attempt %d failed, retry %d/%d"
                    % (task_id, attempt + 1, retry_count,
                       self.max_retries))
                if self.retry_backoff_seconds:
                    time.sleep(self.retry_backoff_seconds)
                continue
            break
        # 超时转 failed：超时是失败，不是取消；最终状态必须是
        # VALID_STATUSES 里的 failed，绝不能把 "timeout" 写进库。
        if result is not None and result.get("status") == "timeout":
            result = dict(result, status="failed")
            error = "task timeout after %s seconds" % self.task_timeout
            self._audit(task["project_id"], "worker_timeout", task_id)

        # 取消优先：执行期间被 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"
            if retry_count:
                error = "[retries=%d] %s" % (retry_count, error)
            # 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。
        """
        # 常驻启动清理兜底：清掉上一轮残留的 codex- 容器，
        # 失败只记日志（见 cleanup_stale_containers），不炸循环。
        self.cleanup_stale_containers()
        # 运维心跳：启动时先写一次，之后每轮结束再写，server 的
        # /health 据此算 worker_ok（见 server.py 的健康检查）。
        self._write_heartbeat()
        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)
                # 每轮结束写心跳：哪怕本轮一个任务都没派，worker 活着
                # 这件事本身也要被 server 看到。
                self._write_heartbeat()
                # 按 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("--heartbeat-file", default=None,
                        help="运维心跳文件路径（默认：环境变量 "
                             "CLOUD_WORKER_HEARTBEAT_FILE，未设则为 "
                             "./worker-heartbeat.json；主机部署时用此"
                             "参数指到 /var/lib/codex-app-server/ 下，"
                             "与 server 的 health 读取路径保持一致）")
    parser.add_argument("--repos-dir", default="./repos",
                        help="项目仓库检出根目录（默认 ./repos）；项目有"
                             " repo_url 时 worker 在此 clone/fetch 各项目"
                             "仓库；Docker 执行器下它必须放在"
                             " --docker-extra-mounts 已挂载的目录下，"
                             "否则容器内 task repo 路径不存在")
    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("--task-timeout", type=float, default=1800,
                        help="每任务截止时间（秒，默认 1800）；超时后真杀"
                             "容器/进程组、任务置 failed；与 --docker-timeout"
                             "（单次 docker run 超时）取较小值，与"
                             " --max-seconds（run_loop 整体墙钟上限）不同")
    parser.add_argument("--max-retries", type=int, default=1,
                        help="模型任务失败重试次数（默认 1，0 表示不重试；"
                             "取消与已记 fallback 的失败不重试）")
    parser.add_argument("--retry-backoff-seconds", type=float, default=5.0,
                        help="重试前等待秒数（默认 5）")
    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),
                    task_timeout=args.task_timeout,
                    max_retries=args.max_retries,
                    retry_backoff_seconds=args.retry_backoff_seconds,
                    repos_dir=args.repos_dir,
                    heartbeat_file=args.heartbeat_file)
    # CLI 启动清理兜底（常驻与单次都做）：清掉上一轮残留容器。
    worker.cleanup_stale_containers()
    # 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()