codex-app-server/examples/harness2_sweep.py

"""多轮 Harness² 改进批量跑（sweep，接线层 CLI）。

读一个 JSON manifest，顺序跑多轮 run_improvement_via_tasks：所有轮
共享同一个 library_root（已入库的经验会被后续轮读到，即上游 sweep
的 library restore 语义）；单轮失败只记入报告、不中断后续轮
（dead 隔离）；最后写 sweep-report.json 并汇总打印。

这会创建真实任务、花真实模型配额——所以默认只打印计划，不建
任务：先跑一遍不带 --yes 看计划，确认无误后，再带 --yes 重跑。
密钥不走命令行参数（ps 可见）：用环境变量 CODEX_API_KEY。

manifest 形状：
{
  "runs": [
    {"run_tag": "sw1", "task_id": "ledger-1",
     "prompt": "Sum every signed amount in ledger.csv.",
     "base_harness_dir": "./base-harness", "source_dir": "./task-source",
     "domain": "accounting", "k": 1, "mode": "parallel",
     "budget": 5, "max_sessions": null, "model": null,
     "role_models": null, "effort": null, "approval": null,
     "environment_id": null, "executor": null,
     "no_time_limit": false, "timeout": 3600}
  ]
}
prompt 也可用 prompt_file（相对 manifest 所在目录）；base_harness_dir
与 source_dir 同样相对 manifest 目录解析（绝对路径原样用）。
"""

import argparse
import json
import os
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from cloud.harness2 import read_tree  # noqa: E402
from cloud.harness2_tasks import (  # noqa: E402
    HttpTaskClient, preflight_checks, run_sweep, validate_run_tag,
)


def _flatten(root: Path) -> dict[str, bytes]:
    out: dict[str, bytes] = {}
    if not root or not root.is_dir():
        return out
    for p in sorted(root.rglob("*")):
        if p.is_file() and not p.is_symlink() and "__pycache__" not in p.parts:
            out[str(p.relative_to(root))] = p.read_bytes()
    return out


def _resolve(manifest_dir: Path, value):
    if not value:
        return value
    p = Path(value)
    return p if p.is_absolute() else manifest_dir / p


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True, help="JSON manifest 路径")
    ap.add_argument("--base-url", required=True, help="控制面地址")
    ap.add_argument("--project-id", required=True)
    ap.add_argument("--output", default="",
                    help="sweep 工作根目录（默认临时目录）")
    ap.add_argument("--library-root", default="",
                    help="共享经验库根目录（空 = 各轮不入库）")
    ap.add_argument("--yes", action="store_true", help="确认创建真实任务并消费配额")
    args = ap.parse_args()

    manifest_path = Path(args.manifest)
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, ValueError) as exc:
        print(f"manifest 读不到或不是合法 JSON：{exc}", file=sys.stderr)
        return 2
    runs = manifest.get("runs")
    if not isinstance(runs, list) or not runs:
        print("manifest 里 runs 必须是非空列表", file=sys.stderr)
        return 2
    manifest_dir = manifest_path.resolve().parent

    specs = []
    seen_tags = set()
    for i, raw in enumerate(runs):
        if not isinstance(raw, dict):
            print(f"runs[{i}] 不是对象", file=sys.stderr)
            return 2
        tag = raw.get("run_tag", "")
        try:
            if not tag:
                raise ValueError("每轮必须给 run_tag")
            validate_run_tag(tag)
        except ValueError as exc:
            print(f"runs[{i}] run_tag 非法：{exc}", file=sys.stderr)
            return 2
        if tag in seen_tags:
            print(f"runs[{i}] run_tag 重复：{tag}", file=sys.stderr)
            return 2
        seen_tags.add(tag)
        prompt = raw.get("prompt", "")
        if raw.get("prompt_file"):
            try:
                prompt = _resolve(manifest_dir, raw["prompt_file"]).read_text(
                    encoding="utf-8")
            except OSError as exc:
                print(f"runs[{i}] prompt_file 读不到：{exc}", file=sys.stderr)
                return 2
        base_dir = _resolve(manifest_dir, raw.get("base_harness_dir", ""))
        errors, warnings = preflight_checks(
            base_harness_dir=base_dir,
            source_dir=_resolve(manifest_dir, raw.get("source_dir")) or None,
            prompt=prompt, run_tag=tag, base_url=args.base_url)
        for e in errors:
            print(f"FAIL runs[{i}] {tag}: {e}", file=sys.stderr)
        for w in warnings:
            print(f"WARN runs[{i}] {tag}: {w}")
        if errors:
            return 2
        source_dir = _resolve(manifest_dir, raw.get("source_dir"))
        specs.append({
            "run_tag": tag,
            "task_id": raw.get("task_id", f"h2-task-{tag}"),
            "prompt": prompt,
            "base_harness": read_tree(base_dir),
            "source_files": _flatten(source_dir) if source_dir else None,
            "domain": raw.get("domain", "general"),
            "k": raw.get("k", 1),
            "mode": raw.get("mode", "parallel"),
            "budget_usd": raw.get("budget"),
            "max_sessions": raw.get("max_sessions"),
            "model": raw.get("model"),
            "role_models": raw.get("role_models"),
            "effort": raw.get("effort"),
            "approval": raw.get("approval"),
            "environment_id": raw.get("environment_id"),
            "executor": raw.get("executor"),
            "no_time_limit": bool(raw.get("no_time_limit", False)),
            "timeout_seconds": raw.get("timeout", 3600),
        })

    print(f"计划：项目 {args.project_id}，共 {len(specs)} 轮，"
          f"tags={sorted(seen_tags)}，"
          f"共享 library={'是' if args.library_root else '否（不入库）'}")
    if not args.yes:
        print("未带 --yes：只打印计划，不创建任何任务。确认无误后加 --yes 重跑。")
        return 0
    api_key = os.environ.get("CODEX_API_KEY")
    if not api_key:
        print("缺少 CODEX_API_KEY 环境变量（密钥不接受命令行参数）", file=sys.stderr)
        return 2

    work_root = Path(args.output) if args.output else Path(
        tempfile.mkdtemp(prefix="h2sweep-"))
    client = HttpTaskClient(args.base_url, api_key)
    report = run_sweep(
        client, project_id=args.project_id, runs=specs, work_root=work_root,
        library_root=Path(args.library_root) if args.library_root else None)
    report_path = work_root / "sweep-report.json"
    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2),
                           encoding="utf-8")
    for entry in report["runs"]:
        status = "ok" if entry["ok"] else f"FAIL {entry['error']}"
        print(f"{entry['run_tag']}: {status} "
              f"chosen={entry['chosen_label']} "
              f"cost={entry['total_cost_usd']}")
    totals = report["totals"]
    print(f"汇总：{totals['n_ok']}/{totals['n_runs']} 轮成功，"
          f"总成本 {totals['total_cost_usd']}，报告 {report_path}")
    return 0 if totals["n_ok"] == totals["n_runs"] else 1


if __name__ == "__main__":
    raise SystemExit(main())