codex-app-server/cloud/tests/test_harness2.py

"""cloud/harness2.py 单测：递归改进循环。

不跑真实模型、不联网：执行器在进程内跑 harness 树里的脚本，
编辑器/选择器/反思器都是确定性假实现（上游 harness2-demo 的同款
做法：base 脚本只加正数，编辑器去掉正数过滤，正确总数是带符号和）。

账本：100.25, -20.10, 5.00 —— 带符号和 85.15，只加正数是 105.25。
"""

import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

try:
    from cloud.harness2 import (
        Library, Rollout, EditResult, validate,
        read_tree, write_tree, is_allowed, normalize_rel, autotrigger,
        diff_summary, diversity, improve_task, select_final,
        assign_ids,
    )
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.harness2 import (
        Library, Rollout, EditResult, validate,
        read_tree, write_tree, is_allowed, normalize_rel, autotrigger,
        diff_summary, diversity, improve_task, select_final,
        assign_ids,
    )

BASE_SCRIPT = """import csv
import sys
from decimal import Decimal
with open(sys.argv[1]) as f:
    amounts = [Decimal(r['amount']) for r in csv.DictReader(f)]
print(f"Total: {sum((a for a in amounts if a > 0), Decimal(0)):.2f}")
"""

BASE_TREE = {
    "AGENTS.md": "Use scripts/summarize.py to total the ledger.\n",
    "scripts/summarize.py": BASE_SCRIPT,
}


class LedgerAdapter:
    """任务：对账本求和。执行 = 真跑 harness 树里的脚本（无模型）。"""

    task_id = "ledger-1"
    domain = "accounting"
    prompt = "Sum every signed amount in ledger.csv. Return Total: with two decimals."

    def __init__(self, root):
        self.root = Path(root)
        self.executions = 0
        self.ledger = self.root / "ledger.csv"
        self.ledger.write_text("amount\n100.25\n-20.10\n5.00\n", encoding="utf-8")

    def base_harness(self):
        return dict(BASE_TREE)

    def execute(self, label, harness):
        self.executions += 1
        work = self.root / "runs" / label
        work.mkdir(parents=True, exist_ok=True)
        write_tree(harness, work / "harness")
        proc = subprocess.run(
            [sys.executable, str(work / "harness" / "scripts" / "summarize.py"),
             str(self.ledger)],
            capture_output=True, text=True, timeout=10,
        )
        out = proc.stdout.strip()
        return Rollout(
            label=label, ok=proc.returncode == 0, delivered=out.startswith("Total:"),
            deliverable_text=out, trajectory="[ACTION] Run scripts/summarize.py\n"
            "[OBSERVATION]\n" + out, cost_usd=0.01,
            error="" if proc.returncode == 0 else proc.stderr[-200:],
        )

    def validate_tree(self, harness):
        if "scripts/summarize.py" not in harness:
            raise ValueError("The local executor requires scripts/summarize.py")


class FixEditor:
    """确定性编辑器：去掉正数过滤；同时在 context/harness/ 里真实改文件，
    模拟 codex 任务编辑工作目录的方式，再把改完的树读回。"""

    def edit(self, mission, context_dir, label):
        harness_dir = Path(context_dir) / "harness"
        script = harness_dir / "scripts" / "summarize.py"
        script.write_text(script.read_text().replace(" if a > 0", ""), encoding="utf-8")
        return EditResult(tree=read_tree(harness_dir),
                          notes="Keep signed amounts: the base total omitted credits.")


class BadEditor:
    """永远交回非法树：AGENTS.md 为空，校验必拒，最终必须回落到 base。"""

    def edit(self, mission, context_dir, label):
        return EditResult(tree={"AGENTS.md": ""}, notes="bad proposal")


class PickSignedSelector:
    def select(self, mission, context_dir):
        for path in sorted((Path(context_dir) / "candidates").glob("*.md")):
            if "Total: 85.15" in path.read_text():
                return path.stem, "The signed total matches the source ledger.", 0.0
        raise ValueError("no candidate with the signed total")


class PickFirstSelector:
    def select(self, mission, context_dir):
        paths = sorted((Path(context_dir) / "candidates").glob("*.md"))
        return paths[0].stem, "first delivered candidate", 0.0


class RaisingSelector:
    def select(self, mission, context_dir):
        raise RuntimeError("selector session failed")


class LedgerReflector:
    def reflect(self, mission, context_dir):
        root = Path(context_dir)
        changed = (root / "harness" / "scripts" / "summarize.py").read_text() != (
            root / "base_harness" / "scripts" / "summarize.py").read_text()
        if not changed:
            return {}
        return {
            "cue": "When a ledger includes credits and debits",
            "lesson": "Include both signs when summing in scripts/summarize.py.",
            "evidence": "Removing the positive-only filter changed the probe total "
                        "to include the credits omitted by the base trajectory.",
            "confidence": "high for this observed arithmetic contrast",
            "edits": "scripts/summarize.py: dropped the positive-only filter; observed: "
                     "probe total included the omitted credits; verdict: helped",
        }


class TestTree(unittest.TestCase):
    def test_validate_accepts_base(self):
        census = validate(dict(BASE_TREE))
        self.assertEqual(census["scripts"], 1)
        self.assertEqual(census["n_files"], 2)

    def test_validate_rejects_skill_without_description(self):
        tree = dict(BASE_TREE)
        tree[".agents/skills/check/SKILL.md"] = "Do the check.\n"
        tree["AGENTS.md"] += "- When checking, read .agents/skills/check/SKILL.md.\n"
        with self.assertRaises(ValueError) as cm:
            validate(tree)
        self.assertIn("description", str(cm.exception))

    def test_validate_rejects_untriggered_skill(self):
        tree = dict(BASE_TREE)
        tree[".agents/skills/check/SKILL.md"] = "---\ndescription: check totals\n---\nDo it.\n"
        with self.assertRaises(ValueError) as cm:
            validate(tree)
        self.assertIn("never named", str(cm.exception))

    def test_validate_rejects_bad_hooks_and_bad_script(self):
        tree = dict(BASE_TREE)
        tree[".codex/hooks.json"] = '{"hooks": {"Nope": []}}'
        with self.assertRaises(ValueError):
            validate(tree)
        tree = dict(BASE_TREE)
        tree["scripts/summarize.py"] = "def broken(:\n"
        with self.assertRaises(ValueError):
            validate(tree)
        tree = dict(BASE_TREE)
        tree["AGENTS.md"] = "x" * 8001
        with self.assertRaises(ValueError):
            validate(tree)

    def test_validate_rejects_missing_script_ref(self):
        tree = dict(BASE_TREE)
        tree["AGENTS.md"] += "Run python3 scripts/nope.py first.\n"
        with self.assertRaises(ValueError):
            validate(tree)

    def test_autotrigger_repairs_and_respects_cap(self):
        tree = dict(BASE_TREE)
        tree[".agents/skills/check/SKILL.md"] = "---\ndescription: check totals\n---\nDo it.\n"
        fixed = autotrigger(tree)
        self.assertEqual(fixed, ["check"])
        self.assertIn("check", tree["AGENTS.md"])
        validate(tree)  # 补完触发行后应能通过
        capped = dict(BASE_TREE)
        capped["AGENTS.md"] = "x" * 7995
        capped[".agents/skills/check/SKILL.md"] = "---\ndescription: check totals\n---\nDo it.\n"
        self.assertEqual(autotrigger(capped), [])  # 超上限时跳过，不硬塞

    def test_paths_and_roundtrip(self):
        with self.assertRaises(ValueError):
            normalize_rel("../escape.py")
        self.assertFalse(is_allowed("/etc/passwd"))
        self.assertFalse(is_allowed("evil.py"))
        self.assertTrue(is_allowed(".agents/skills/a/SKILL.md"))
        with tempfile.TemporaryDirectory() as tmp:
            dest = Path(tmp) / "tree"
            write_tree(dict(BASE_TREE), dest)
            self.assertEqual(read_tree(dest), BASE_TREE)
            rejected = []
            (dest / "evil.py").write_text("x", encoding="utf-8")
            self.assertEqual(read_tree(dest, rejected=rejected), BASE_TREE)
            self.assertEqual(rejected, ["evil.py"])

    def test_diff_and_diversity(self):
        changed = dict(BASE_TREE)
        changed["scripts/summarize.py"] = "changed"
        d = diff_summary(BASE_TREE, changed)
        self.assertEqual(d["modified"], ["scripts/summarize.py"])
        self.assertEqual(d["added"], [])
        div = diversity({"a": BASE_TREE, "b": changed})
        self.assertEqual(div["n_distinct"], 2)
        self.assertLess(div["mean_line_jaccard"], 1.0)


class TestImprove(unittest.TestCase):
    def _run(self, mode, editor=None, selector=None, budget=None, k=1):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        adapter = LedgerAdapter(root)
        library = Library(root / "library", "accounting")
        result = improve_task(
            adapter, editor or FixEditor(),
            selector=selector or PickSignedSelector(),
            reflector=LedgerReflector(), work_dir=root / "work",
            k=k, mode=mode, library=library, budget_usd=budget)
        return result, adapter, library, root

    def test_parallel_k1_improves_and_enters_library(self):
        result, adapter, library, root = self._run("parallel")
        self.assertTrue(result.ok)
        self.assertEqual(result.selection.chosen_label, "agg1")
        chosen = result.mode_result.rollouts["agg1"]
        self.assertEqual(chosen.deliverable_text, "Total: 85.15")
        self.assertEqual(adapter.executions, 4)  # base + inv1 + dom1 + agg1
        self.assertIsNotNone(result.library_entry)
        self.assertIn("ledger-1", library.entries())
        index = (root / "library" / "accounting" / "INDEX.md").read_text()
        self.assertIn("credits and debits", index)
        self.assertIn("85.15", result.selection.why + chosen.deliverable_text)

    def test_sequential_k1_improves(self):
        result, adapter, library, root = self._run("sequential")
        self.assertTrue(result.ok)
        self.assertEqual(result.selection.chosen_label, "c1")
        self.assertEqual(result.mode_result.rollouts["c1"].deliverable_text, "Total: 85.15")
        self.assertEqual(adapter.executions, 4)  # base + inv1 + dom1 + c1

    def test_invalid_proposal_falls_back_to_base(self):
        result, adapter, library, root = self._run(
            "parallel", editor=BadEditor(), selector=PickFirstSelector())
        self.assertTrue(result.ok)  # 回落的 base 产物仍可交付
        self.assertEqual(set(result.mode_result.fell_back), {"inv1", "dom1", "agg1"})
        self.assertEqual(result.mode_result.trees["agg1"], BASE_TREE)
        self.assertIsNone(result.library_entry)  # 树没变，反思无经验，不入库

    def test_budget_exceeded_blocks_library_but_still_delivers(self):
        result, adapter, library, root = self._run("parallel", budget=0.03)
        self.assertTrue(result.ok)
        self.assertTrue(result.budget_exceeded)
        self.assertIsNone(result.library_entry)

    def test_k_validation(self):
        with tempfile.TemporaryDirectory() as tmp:
            adapter = LedgerAdapter(Path(tmp))
            with self.assertRaises(ValueError):
                improve_task(adapter, FixEditor(), work_dir=Path(tmp) / "w", k=0)
            with self.assertRaises(ValueError):
                improve_task(adapter, FixEditor(), work_dir=Path(tmp) / "w", k=4)
            with self.assertRaises(ValueError):
                improve_task(adapter, FixEditor(), work_dir=Path(tmp) / "w", k=1,
                             mode="sideways")

    def test_trivial_k1_without_selector(self):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        adapter = LedgerAdapter(root)
        result = improve_task(adapter, FixEditor(), selector=None,
                              reflector=LedgerReflector(), work_dir=root / "work",
                              k=1, mode="parallel")
        self.assertTrue(result.ok)
        self.assertTrue(result.selection.trivial)


class TestSelection(unittest.TestCase):
    def test_drops_undelivered_and_falls_back_on_selector_error(self):
        with tempfile.TemporaryDirectory() as tmp:
            adapter = LedgerAdapter(Path(tmp))
            rollouts = {
                "agg1": Rollout(label="agg1", ok=True, delivered=True,
                                deliverable_text="Total: 1.00"),
                "agg2": Rollout(label="agg2", ok=False, delivered=False),
            }
            # k=2 的终版是 agg1/agg2；agg2 未交付应被丢掉，agg1 是唯一候选
            sel = select_final("ledger-1", 2, "parallel", rollouts, adapter=adapter,
                               selector=RaisingSelector(), work_dir=Path(tmp) / "work",
                               base_tree=dict(BASE_TREE))
            self.assertTrue(sel.ok)
            self.assertTrue(sel.fallback)
            self.assertEqual(sel.chosen_label, "agg1")
            self.assertIn("agg2", sel.dropped)

    def test_assign_ids_is_stable(self):
        a = assign_ids(["agg2", "agg1"], task_id="t", k=1, mode="parallel")
        b = assign_ids(["agg1", "agg2"], task_id="t", k=1, mode="parallel")
        self.assertEqual(a, b)  # 输入顺序不影响匿名映射


class TestLibrary(unittest.TestCase):
    def test_add_requires_experience(self):
        with tempfile.TemporaryDirectory() as tmp:
            lib = Library(Path(tmp) / "lib", "accounting")
            self.assertIsNone(lib.add("t1", dict(BASE_TREE), experience={}))
            self.assertEqual(lib.entries(), [])
            self.assertIsNone(lib.context_dir())
            entry = lib.add("t1", dict(BASE_TREE), why="why", label="agg1",
                            experience={"cue": "c", "lesson": "l", "evidence": "e"})
            self.assertIsNotNone(entry)
            self.assertEqual(lib.entries(), ["t1"])
            self.assertEqual(lib.tree_of("t1"), BASE_TREE)
            self.assertIsNotNone(lib.context_dir())


if __name__ == "__main__":
    unittest.main()