codex-app-server/cloud/worker_sweep.py

"""管理员 CLI：失联 worker 清扫（workers kill 接线，A5）。

用法：python3 -m cloud.worker_sweep --db <控制面库路径> [--dry-run]
    [--stale-seconds N] [--host <本机主机名>]
    API key 走环境变量 CLOUD_WORKER_SWEEP_API_KEY，或 --api-key
    显式给（命令行会进进程表，生产优先用环境变量）。

这是 reap_stale_workers 的第一个生产调用方（此前除 workers.py
本体与测试外全仓库零生产调用方）：

- 鉴权：先经 AuthStore.authenticate 验 key，再认
  CODEX_GLOBAL_ADMINS（与 handler_admin._is_global_admin 同款
  口径：逗号分隔的用户 id 清单，未配置或为空时一律 fail-closed，
  绝不回落到项目角色推导）。任一步失败都抛 PermissionError，
  且在调 reap 之前就拒掉——未授权调用绝不走到判定/杀。
  被拒的已认证非管理员用户记一条 admin_worker_sweep_denied
  审计（审计写失败不挡拒绝本身）；key 本身验不过时没有可
  归属的用户，不记这条。
- 杀前审计 + 执行后复核：授权后直接调 reap_stale_workers，
  复用 classify_stale_workers 的判定，不另造判定函数；每行
  kill 档在 reap 内先再验 pid 已死、杀前写审计（审计失败则
  不杀，fail-closed）、只给该 pid 本人发信号、杀后用同一
  checker 复核并写杀后审计；本 CLI 另写一条
  admin_worker_sweep 汇总审计（killed / listed / 失败计数，
  写失败只记 audit_summary=False，不改已执行结果）。
- 不新增 HTTP 路由（v1.1 契约确认 B8 未定前不扩 API 面）；
  不枚举容器、不按名/前缀匹配、不碰基础设施容器、不删
  workers 注册行：杀的是判定为 kill 的行对应的 pid 本人，
  注册行的去留是另一条决策，本 CLI 不擅自删。
- --dry-run 只调 classify_stale_workers 列判定，不调 reap、
  不杀、不写杀审计，供运维先看再动手。

本模块只 import cloud.auth 与 cloud.workers，不 import
cloud.server / handler_*，避免循环 import；PG 库路径同样经
cloud.db 的 connect 适配（与 workers 本体同款）。
"""

import argparse
import json
import os
import sys

from cloud.auth import AuthStore
from cloud.workers import classify_stale_workers
from cloud.workers import reap_stale_workers


def global_admin_ids(admin_ids=None):
    """全局管理员清单：显式传入优先，否则读 CODEX_GLOBAL_ADMINS。

    未配置或为空时返回空集——调用方一律 fail-closed。
    """
    if admin_ids is not None:
        return {str(item).strip() for item in admin_ids if str(item).strip()}
    raw = os.environ.get("CODEX_GLOBAL_ADMINS") or ""
    return {part.strip() for part in raw.split(",") if part.strip()}


def sweep_stale_workers_admin(db_path, api_key, stale_seconds=None,
                              local_host=None, pid_is_dead_fn=None,
                              kill_fn=None, audit_fn=None, dry_run=False,
                              admin_ids=None, now=None):
    """管理员清扫入口：鉴权通过后才调 reap（或 dry-run 只列判定）。

    返回 dict：user_id、dry_run、results（每行判定 + 执行字段，
    与 reap/classify 的返回同款）、killed / listed / failed
    三个 worker_id 清单、audit_summary（汇总审计是否写成功，
    dry-run 时为 None）。未授权时抛 PermissionError 且绝不调
    reap、不调 kill_fn。kill_fn / audit_fn / pid_is_dead_fn 仅
    供测试注入；生产调用一律走默认（真 SIGKILL + 统一审计口）。
    """
    store = AuthStore(db_path)
    try:
        user_id = store.authenticate(api_key) if api_key else None
        if user_id is None:
            raise PermissionError(
                "worker sweep requires a valid admin API key")
        if user_id not in global_admin_ids(admin_ids):
            try:
                store.audit(user_id, "", "admin_worker_sweep_denied",
                            "worker sweep forbidden: not a global admin")
            except Exception:
                pass
            raise PermissionError(
                "worker sweep forbidden: not a global admin")
        if dry_run:
            results = classify_stale_workers(
                db_path, stale_seconds=stale_seconds,
                local_host=local_host, pid_is_dead_fn=pid_is_dead_fn,
                now=now)
            return {
                "user_id": user_id,
                "dry_run": True,
                "results": results,
                "killed": [],
                "listed": [row["worker_id"] for row in results],
                "failed": [],
                "audit_summary": None,
            }
        results = reap_stale_workers(
            db_path, stale_seconds=stale_seconds, local_host=local_host,
            pid_is_dead_fn=pid_is_dead_fn, kill_fn=kill_fn,
            audit_fn=audit_fn, now=now)
        killed = [row["worker_id"] for row in results
                  if row.get("outcome") == "killed"]
        listed = [row["worker_id"] for row in results
                  if row.get("outcome") == "listed"]
        failed = [row["worker_id"] for row in results
                  if row.get("outcome") in (
                      "kill_failed_pid_alive", "audit_failed")]
        detail = ("killed=%d listed=%d failed=%d killed_ids=%s"
                  % (len(killed), len(listed), len(failed),
                     ",".join(killed)))
        try:
            store.audit(user_id, "", "admin_worker_sweep", detail)
            audit_summary = True
        except Exception:
            audit_summary = False
        return {
            "user_id": user_id,
            "dry_run": False,
            "results": results,
            "killed": killed,
            "listed": listed,
            "failed": failed,
            "audit_summary": audit_summary,
        }
    finally:
        try:
            store.conn.close()
        except Exception:
            pass


def main(argv=None):
    """CLI 入口：鉴权失败 exit 1，参数缺失 exit 2，成功 exit 0。

    成功时 stdout 只打 JSON 汇总（worker_id 与 outcome，不含
    任何密钥）；失败原因打 stderr，不回 JSON。
    """
    parser = argparse.ArgumentParser(
        description="管理员清扫失联 worker（只杀 pid 已死验证过的本机行）")
    parser.add_argument("--db", required=True, help="控制面数据库路径")
    parser.add_argument("--api-key", default=None,
                        help="管理员 API key；不给时读环境变量 "
                             "CLOUD_WORKER_SWEEP_API_KEY")
    parser.add_argument("--stale-seconds", type=float, default=None,
                        help="失联阈值秒数；不给时复用 alive 阈值语义")
    parser.add_argument("--host", default=None,
                        help="本机主机名；不给时取 socket.gethostname()")
    parser.add_argument("--dry-run", action="store_true",
                        help="只列判定，不杀")
    args = parser.parse_args(argv)
    api_key = args.api_key or os.environ.get("CLOUD_WORKER_SWEEP_API_KEY")
    try:
        summary = sweep_stale_workers_admin(
            args.db, api_key, stale_seconds=args.stale_seconds,
            local_host=args.host, dry_run=args.dry_run)
    except PermissionError as exc:
        print("worker sweep refused: %s" % exc, file=sys.stderr)
        raise SystemExit(1)
    print(json.dumps({
        "user_id": summary["user_id"],
        "dry_run": summary["dry_run"],
        "killed": summary["killed"],
        "listed": summary["listed"],
        "failed": summary["failed"],
        "audit_summary": summary["audit_summary"],
    }, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    main()