"""cloud/harness2_tasks.py 单测：真实 codex 任务版接线。

不跑真实模型、不联网：FakeTaskClient 在进程内模拟控制面——探针
任务真跑 bundle 里的脚本，编辑任务按文件块格式回改动，选择与
反思按最终回复格式回结果。HTTP 客户端则对本地假服务端实测。
账本与 base 树同 test_harness2.py：带符号和 85.15，只加正数 105.25。
"""

import csv
import json
import subprocess
import sys
import tempfile
import threading
import unittest
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

try:
    from cloud import harness2_tasks
    from cloud.harness2 import (
        parse_choice, parse_deliverable_files, parse_experience,
        parse_harness_files, parse_no_change, parse_notes, parse_why,
    )
    from cloud.harness2_tasks import (
        CodexEditor, CodexProbeAdapter, HttpTaskClient, TaskState,
        build_multipart, preflight_checks, run_improvement_via_tasks, run_sweep,
        run_task, unwrap_result_text, validate_run_tag,
    )
    from cloud.uploads import parse_multipart, parse_paths_field
except ImportError:  # discover -s cloud/tests 时把项目根目录补进 sys.path
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud import harness2_tasks
    from cloud.harness2 import (
        parse_choice, parse_deliverable_files, parse_experience,
        parse_harness_files, parse_no_change, parse_notes, parse_why,
    )
    from cloud.harness2_tasks import (
        CodexEditor, CodexProbeAdapter, HttpTaskClient, TaskState,
        build_multipart, preflight_checks, run_improvement_via_tasks, run_sweep,
        run_task, unwrap_result_text, validate_run_tag,
    )
    from cloud.uploads import parse_multipart, parse_paths_field

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,
}

LEDGER = b"amount\n100.25\n-20.10\n5.00\n"


class FakeTaskClient:
    """进程内控制面：按 prompt 类型分派到确定性假 worker。"""

    def __init__(self, editor_mode="fix", selector_mode="signed", envelope=False,
                 fail_tags=(), endpoint_files=None, endpoint_error=False,
                 probe_extra=""):
        self.editor_mode = editor_mode
        self.selector_mode = selector_mode
        self.envelope = envelope
        self.fail_tags = tuple(fail_tags)
        self.endpoint_files = dict(endpoint_files or {})
        self.endpoint_error = endpoint_error
        self.probe_extra = probe_extra
        self.uploads = {}
        self.tasks = {}
        self.by_key = {}
        self.created = 0
        self.cancelled = []
        self.task_models = {}

    def upload(self, files):
        uid = f"u{len(self.uploads) + 1}"
        self.uploads[uid] = dict(files)
        return uid

    def create_task(self, *, project_id, prompt, upload_id=None,
                    idempotency_key=None, **kwargs):
        if idempotency_key and any(t in idempotency_key for t in self.fail_tags):
            raise RuntimeError("injected task failure")
        if idempotency_key and idempotency_key in self.by_key:
            return self.by_key[idempotency_key]
        tid = f"t{len(self.tasks) + 1}"
        self.created += 1
        files = self.uploads.get(upload_id, {}) if upload_id else {}
        result = self._work(prompt, files)
        if self.envelope:
            # 模拟控制面真实形状：result_text 是 JSON 信封，text 字段
            # 里才是最终回复——接线层必须拆信封，否则行首锚定的文件
            # 块解析永远看不见块（2026-10-29 真实轮次根因）。
            result = json.dumps({
                "task": tid, "thread_id": "th", "cwd": "/x",
                "status": "completed", "text": result, "error": None,
            })
        self.tasks[tid] = TaskState(task_id=tid, status="completed",
                                    result_text=result)
        self.task_models[tid] = kwargs.get("model")
        if idempotency_key:
            self.by_key[idempotency_key] = tid
        return tid

    def _work(self, prompt, files):
        if prompt.startswith("You are editing an agent harness"):
            if self.editor_mode == "gibberish":
                return "hello, I did stuff"
            if self.editor_mode in ("notes-only-then-fix", "notes-only"):
                if ("YOUR PREVIOUS REPLY WAS REJECTED" in prompt
                        and self.editor_mode == "notes-only-then-fix"):
                    script = files["context/harness/scripts/summarize.py"].decode()
                    fixed = script.replace(" if a > 0", "")
                    return (
                        "=== HARNESS FILE: scripts/summarize.py ===\n" + fixed +
                        "=== END HARNESS FILE ===\n"
                        "<notes>Removed the positive-only filter.</notes>")
                # 真实轮次原样复现：notes 里声称已修正，却不发文件块。
                return ("<notes>Removed the positive-only filter so signed "
                        "amounts are summed; the script now totals all amounts."
                        "</notes>")
            if self.editor_mode == "no-change":
                return ("<notes>No evidence justifies a change.</notes>"
                        "<no-change>The base trajectory shows no defect."
                        "</no-change>")
            script = files["context/harness/scripts/summarize.py"].decode()
            fixed = script.replace(" if a > 0", "")
            return (
                "=== HARNESS FILE: scripts/summarize.py ===\n" + fixed +
                "=== END HARNESS FILE ===\n"
                "<notes>Keep signed amounts: the base total omitted credits.</notes>"
            )
        if prompt.startswith("You are choosing which"):
            if self.selector_mode == "gibberish":
                return "I like them all"
            cands = sorted(p for p in files if p.startswith("context/candidates/")
                           and p.endswith(".md"))
            for p in cands:
                if b"Total: 85.15" in files[p]:
                    stem = Path(p).stem
                    return (f"```choice\n{stem}\n```\n"
                            "<why>The signed total matches the source ledger.</why>")
            stem = Path(cands[0]).stem
            return (f"```choice\n{stem}\n```\n<why>first delivered candidate</why>")
        if prompt.startswith("You are reading the complete record"):
            changed = (files.get("context/harness/scripts/summarize.py")
                       != files.get("context/base_harness/scripts/summarize.py"))
            if not changed:
                return "NO_ENTRY"
            return (
                "<edits>\nscripts/summarize.py: dropped the positive-only filter; "
                "verdict: helped\n</edits>\n"
                "<experience>\ncue: When a ledger includes credits and debits\n"
                "lesson: Include both signs when summing in scripts/summarize.py.\n"
                "evidence: Removing the positive-only filter changed the probe total "
                "to include the credits omitted by the base trajectory.\n"
                "confidence: high\n</experience>"
            )
        # 探针：真跑 bundle 里的脚本对 bundle 里的账本求和。
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            for rel, data in files.items():
                dest = root / rel
                dest.parent.mkdir(parents=True, exist_ok=True)
                dest.write_bytes(data)
            proc = subprocess.run(
                [sys.executable, str(root / "scripts" / "summarize.py"),
                 str(root / "ledger.csv")],
                capture_output=True, text=True, timeout=10)
            return proc.stdout.strip() + self.probe_extra

    def get_task(self, task_id):
        return self.tasks[task_id]

    def get_task_files(self, task_id, path=""):
        if self.endpoint_error:
            raise RuntimeError("409 files unavailable")
        if not path:
            return {"entries": [
                {"path": p, "size": len(c.encode("utf-8"))}
                for p, c in sorted(self.endpoint_files.items())]}
        if path in self.endpoint_files:
            return {"content": self.endpoint_files[path]}
        raise RuntimeError("no such file")

    def get_trace_text(self, task_id):
        return "[ACTION] fake probe\n[OBSERVATION]\n" + self.tasks[task_id].result_text

    def get_cost_usd(self, task_id):
        return 0.02

    def cancel_task(self, task_id):
        self.cancelled.append(task_id)


class TestParseHelpers(unittest.TestCase):
    def test_multipart_roundtrips_through_server_parser(self):
        files = {"context/task.md": b"do it", "AGENTS.md": b"be nice\n"}
        body, content_type = build_multipart(files)
        fields, parsed = parse_multipart(body, content_type)
        paths = parse_paths_field(fields, parsed)
        self.assertEqual(paths, sorted(files))
        self.assertEqual(dict(parsed), files)

    def test_choice_why_notes_experience(self):
        self.assertEqual(parse_choice("```choice\nB\n```"), "B")
        self.assertIsNone(parse_choice("no choice here"))
        self.assertEqual(parse_why("<why>because</why>"), "because")
        self.assertEqual(parse_notes("<notes>did x</notes>"), "did x")

    def test_parse_deliverable_files_lenient(self):
        text = ("=== DELIVERABLE FILE: a.md ===\nhello\n"
                "=== END DELIVERABLE FILE ===\n"
                "=== DELIVERABLE FILE: ../evil ===\nx\n"
                "=== END DELIVERABLE FILE ===\n"
                "=== DELIVERABLE FILE: b.md ===\nunterminated\n")
        self.assertEqual(parse_deliverable_files(text), {"a.md": "hello\n"})
        many = "".join(
            f"=== DELIVERABLE FILE: f{i}.md ===\nx\n"
            "=== END DELIVERABLE FILE ===\n" for i in range(10))
        self.assertEqual(len(parse_deliverable_files(many)), 8)
        big = ("=== DELIVERABLE FILE: big.md ===\n" + "x" * 70000 + "\n"
               "=== END DELIVERABLE FILE ===\n")
        content = parse_deliverable_files(big)["big.md"]
        self.assertTrue(content.endswith("...[truncated]"))
        exp = parse_experience(
            "<experience>\ncue: c\nlesson: l\nevidence: e\nconfidence: high\n</experience>")
        self.assertEqual((exp["cue"], exp["lesson"], exp["evidence"]), ("c", "l", "e"))
        self.assertEqual(parse_experience("NO_ENTRY"), {})
        self.assertEqual(parse_experience("<experience>\ncue: c\n</experience>"), {})

    def test_harness_file_blocks(self):
        text = ("=== HARNESS FILE: AGENTS.md ===\nhello\n"
                "=== END HARNESS FILE ===\n<notes>n</notes>")
        self.assertEqual(parse_harness_files(text), {"AGENTS.md": "hello\n"})
        self.assertEqual(parse_harness_files("nothing here"), {})
        with self.assertRaises(ValueError):
            parse_harness_files("=== HARNESS FILE: /etc/passwd ===\nx\n"
                                "=== END HARNESS FILE ===")
        with self.assertRaises(ValueError):
            parse_harness_files("=== HARNESS FILE: AGENTS.md ===\nunterminated\n")


class TestFullRunViaTasks(unittest.TestCase):
    def _run(self, client, run_tag="run1", work_sub="work", k=1, max_sessions=None):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        result = run_improvement_via_tasks(
            client, project_id="p1", task_id="ledger-1",
            prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
            base_harness=dict(BASE_TREE), run_tag=run_tag,
            work_dir=root / work_sub, domain="accounting",
            source_files={"ledger.csv": LEDGER},
            k=k, mode="parallel", library_root=root / "library",
            max_sessions=max_sessions)
        return result, client, root

    def test_parallel_k1_via_real_task_shape(self):
        client = FakeTaskClient()
        result, _, root = self._run(client)
        self.assertTrue(result.ok)
        self.assertEqual(result.selection.chosen_label, "agg1")
        self.assertTrue(result.selection.trivial)  # k=1 单终版不开选择会话
        self.assertEqual(result.mode_result.rollouts["agg1"].deliverable_text,
                         "Total: 85.15")
        # 4 探针 + 3 编辑 + 1 反思 = 8 个任务（选择 trivial，不建任务）
        self.assertEqual(client.created, 8)
        self.assertIsNotNone(result.library_entry)
        self.assertTrue((root / "library" / "accounting" / "INDEX.md").is_file())

    def _models_by_role(self, client):
        roles = {}
        for key, tid in client.by_key.items():
            for role in ("probe", "edit", "select", "reflect"):
                if f"-{role}-" in key or key.endswith(f"-{role}"):
                    roles.setdefault(role, set()).add(client.task_models[tid])
        return roles

    def test_role_models_land_on_each_role(self):
        client = FakeTaskClient()
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        result = run_improvement_via_tasks(
            client, project_id="p1", task_id="ledger-1",
            prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
            base_harness=dict(BASE_TREE), run_tag="roles1",
            work_dir=root / "work", domain="accounting",
            source_files={"ledger.csv": LEDGER},
            k=2, mode="parallel", library_root=root / "library",
            model="shared-model",
            role_models={"probe": "probe-model", "editor": "editor-model",
                         "selector": "selector-model",
                         "reflector": "reflector-model"})
        self.assertTrue(result.ok)
        roles = self._models_by_role(client)
        self.assertEqual(roles["probe"], {"probe-model"})
        self.assertEqual(roles["edit"], {"editor-model"})
        self.assertEqual(roles["select"], {"selector-model"})
        self.assertEqual(roles["reflect"], {"reflector-model"})

    def test_role_model_falls_back_to_shared_then_default(self):
        client = FakeTaskClient()
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        result = run_improvement_via_tasks(
            client, project_id="p1", task_id="ledger-1",
            prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
            base_harness=dict(BASE_TREE), run_tag="roles2",
            work_dir=root / "work", domain="accounting",
            source_files={"ledger.csv": LEDGER},
            k=1, mode="parallel", library_root=root / "library",
            model="shared-model",
            role_models={"editor": "editor-model"})
        self.assertTrue(result.ok)
        roles = self._models_by_role(client)
        self.assertEqual(roles["probe"], {"shared-model"})
        self.assertEqual(roles["edit"], {"editor-model"})
        self.assertEqual(roles["reflect"], {"shared-model"})
        # 两者都空时不传模型（服务端默认），假客户端记为 None。
        client2 = FakeTaskClient()
        result2, _, _ = self._run(client2, run_tag="roles3")
        self.assertTrue(result2.ok)
        self.assertEqual(set(client2.task_models.values()), {None})

    def test_probe_deliverable_files_from_endpoint(self):
        client = FakeTaskClient(endpoint_files={
            "report.md": "# R\n",
            "ledger.csv": "amount\n999.00\n",
            "AGENTS.md": BASE_TREE["AGENTS.md"],
        })
        result, _, root = self._run(client, run_tag="assets1")
        self.assertTrue(result.ok)
        agg = result.mode_result.rollouts["agg1"]
        # 新增文件与改动过的 seed 文件进产物；内容与 seed 一致的不进。
        self.assertEqual(agg.deliverable_files,
                         {"report.md": "# R\n",
                          "ledger.csv": "amount\n999.00\n"})
        found = [p for p in root.rglob("report.md")
                 if "deliverables/agg1/files" in str(p)]
        self.assertTrue(found)
        self.assertEqual(found[0].read_text(encoding="utf-8"), "# R\n")

    def test_probe_deliverable_files_from_reply_blocks(self):
        block = ("\n=== DELIVERABLE FILE: page.html ===\n<h1>hi</h1>\n"
                 "=== END DELIVERABLE FILE ===\n")
        client = FakeTaskClient(endpoint_error=True, probe_extra=block)
        result, _, _ = self._run(client, run_tag="assets2")
        self.assertTrue(result.ok)
        self.assertEqual(
            result.mode_result.rollouts["agg1"].deliverable_files,
            {"page.html": "<h1>hi</h1>\n"})

    def test_probe_deliverable_files_endpoint_wins_on_collision(self):
        block = ("\n=== DELIVERABLE FILE: page.html ===\nfrom-reply\n"
                 "=== END DELIVERABLE FILE ===\n")
        client = FakeTaskClient(
            endpoint_files={"page.html": "from-endpoint\n"},
            probe_extra=block)
        result, _, _ = self._run(client, run_tag="assets3")
        self.assertTrue(result.ok)
        self.assertEqual(
            result.mode_result.rollouts["agg1"].deliverable_files,
            {"page.html": "from-endpoint\n"})

    def test_probe_deliverable_files_endpoint_error_is_not_fatal(self):
        client = FakeTaskClient(endpoint_error=True)
        result, _, _ = self._run(client, run_tag="assets4")
        self.assertTrue(result.ok)
        self.assertEqual(
            result.mode_result.rollouts["agg1"].deliverable_files, {})

    def test_same_run_tag_replays_without_new_tasks(self):
        client = FakeTaskClient()
        result1, _, _ = self._run(client, work_sub="work1")
        self.assertTrue(result1.ok)
        created_after_first = client.created
        result2, _, _ = self._run(client, work_sub="work2")
        self.assertTrue(result2.ok)
        self.assertEqual(client.created, created_after_first)  # 幂等键命中，零新建

    def test_editor_gibberish_falls_back(self):
        client = FakeTaskClient(editor_mode="gibberish")
        result, _, _ = self._run(client)
        self.assertTrue(result.ok)  # 回落的 base 产物仍可交付
        self.assertEqual(set(result.mode_result.fell_back), {"inv1", "dom1", "agg1"})
        self.assertIsNone(result.library_entry)

    def test_selector_gibberish_falls_back(self):
        # k=1 时选择 trivial，不走选择器；用 k=2 让选择器真正跑起来。
        client = FakeTaskClient(selector_mode="gibberish")
        result, _, _ = self._run(client, k=2)
        self.assertTrue(result.ok)
        self.assertTrue(result.selection.fallback)


class TestEditorFormatRetry(unittest.TestCase):
    """编辑格式纠正：notes-only 先重试再落盘；显式无改动一次通过；
    两次都 notes-only 则如实抛错（阶段回落，不静默采纳）。"""

    def _editor(self, mode):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        ctx = Path(tmp.name)
        (ctx / "harness" / "scripts").mkdir(parents=True)
        for rel, content in BASE_TREE.items():
            (ctx / "harness" / rel).write_text(content)
        client = FakeTaskClient(editor_mode=mode)
        editor = CodexEditor(client, project_id="p", run_tag="fmt")
        return editor, ctx, client

    def test_notes_only_then_fix_lands(self):
        editor, ctx, client = self._editor("notes-only-then-fix")
        result = editor.edit("You are editing an agent harness", ctx, "inv1.a1")
        self.assertNotIn("if a > 0", result.tree["scripts/summarize.py"])
        self.assertEqual(client.created, 2)  # 首发 + 一次格式纠正

    def test_no_change_marker_keeps_base(self):
        editor, ctx, client = self._editor("no-change")
        result = editor.edit("You are editing an agent harness", ctx, "inv1.a1")
        self.assertEqual(result.tree["scripts/summarize.py"], BASE_SCRIPT)
        self.assertEqual(client.created, 1)

    def test_notes_only_twice_raises(self):
        editor, ctx, client = self._editor("notes-only")
        with self.assertRaises(ValueError):
            editor.edit("You are editing an agent harness", ctx, "inv1.a1")
        self.assertEqual(client.created, 2)

    def test_parse_no_change(self):
        self.assertEqual(parse_no_change("<no-change>no defect</no-change>"),
                         "no defect")
        self.assertEqual(parse_no_change("just notes"), "")


class TestUnwrapResultText(unittest.TestCase):
    """拆信封：信封取 text；其它原样返回，不丢内容。"""

    def test_envelope_unwrapped(self):
        raw = json.dumps({"task": "t1", "status": "completed",
                          "text": "=== HARNESS FILE: a.md ===\nhi\n=== END HARNESS FILE ===\n"})
        self.assertTrue(unwrap_result_text(raw).startswith("=== HARNESS FILE"))

    def test_plain_text_unchanged(self):
        self.assertEqual(unwrap_result_text("Total: 85.15"), "Total: 85.15")
        self.assertEqual(unwrap_result_text('{"no": "text field"}'),
                         '{"no": "text field"}')
        self.assertEqual(unwrap_result_text("{not json"), "{not json")
        self.assertEqual(unwrap_result_text(""), "")


class TestEnvelopeFullRound(unittest.TestCase):
    """信封形状整轮：假控制面按真实信封返回时，编辑块仍能落盘、
    探针交付物是干净文本（而非信封 JSON），经验正常入库。"""

    def test_envelope_round(self):
        client = FakeTaskClient(envelope=True)
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        result = run_improvement_via_tasks(
            client, project_id="p1", task_id="ledger-1",
            prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
            base_harness=dict(BASE_TREE), run_tag="env",
            work_dir=root / "work", domain="accounting",
            source_files={"ledger.csv": LEDGER},
            k=1, mode="parallel", library_root=root / "library")
        self.assertTrue(result.ok)
        self.assertEqual(result.mode_result.rollouts["agg1"].deliverable_text,
                         "Total: 85.15")
        self.assertNotIn("thread_id",
                         result.mode_result.rollouts["agg1"].deliverable_text)
        self.assertIsNotNone(result.library_entry)


class TestSessionBudget(unittest.TestCase):
    """会话数门：0 挡全部、2 挡在中途；fmt2 与选择/反思不计入。"""

    def _run(self, max_sessions):
        client = FakeTaskClient()
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        result = run_improvement_via_tasks(
            client, project_id="p1", task_id="ledger-1",
            prompt="Sum every signed amount in ledger.csv. Return Total: with two decimals.",
            base_harness=dict(BASE_TREE), run_tag="sess",
            work_dir=root / "work", domain="accounting",
            source_files={"ledger.csv": LEDGER},
            k=1, mode="parallel", library_root=root / "library",
            max_sessions=max_sessions)
        return result, client

    def test_zero_blocks_everything(self):
        result, client = self._run(0)
        self.assertEqual(client.created, 0)
        self.assertEqual(set(result.mode_result.fell_back), {"inv1", "dom1", "agg1"})

    def test_mid_cap_stops_after_two_sessions(self):
        # base 探针（1）+ inv1 编辑（2），之后 dom1/agg1 与探针全被门挡住。
        result, client = self._run(2)
        self.assertEqual(client.created, 2)
        self.assertIn("dom1", result.mode_result.fell_back)
        self.assertIn("agg1", result.mode_result.fell_back)
        self.assertNotIn("inv1", result.mode_result.fell_back)


class TestValidateRunTag(unittest.TestCase):
    def test_valid(self):
        for tag in ("h2-20261029-real1", "a", "a.b_c-1", "run_1"):
            validate_run_tag(tag)

    def test_invalid(self):
        for tag in ("-bad", "has/slash", "has space", "s1-foo", "s12"):
            with self.assertRaises(ValueError, msg=tag):
                validate_run_tag(tag)
        validate_run_tag("")  # 空串表示稍后生成，不校验


class TestPreflight(unittest.TestCase):
    def _dirs(self, agents="Use scripts/summarize.py.\n"):
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        (root / "base").mkdir()
        if agents is not None:
            (root / "base" / "AGENTS.md").write_text(agents)
        (root / "source").mkdir()
        (root / "source" / "ledger.csv").write_bytes(LEDGER)
        return root

    def test_ok(self):
        root = self._dirs()
        errors, _ = preflight_checks(
            base_harness_dir=root / "base", source_dir=root / "source",
            prompt="do it", run_tag="h2-ok")
        self.assertEqual(errors, [])

    def test_errors(self):
        root = self._dirs(agents="")
        errors, _ = preflight_checks(
            base_harness_dir=root / "base", source_dir=root / "missing",
            prompt="  ", run_tag="s1-bad")
        self.assertEqual(len(errors), 4)  # AGENTS 为空、source 缺失、提示为空、tag 非法
        errors, _ = preflight_checks(base_harness_dir=root / "nope", prompt="x")
        self.assertTrue(any("不存在" in e for e in errors))


class TestHttpClient(unittest.TestCase):
    def test_client_against_local_server(self):
        seen = {"keys": [], "uploads": [], "tasks": []}
        polls = {"t1": 0}

        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass

            def _json(self, status, obj):
                raw = json.dumps(obj).encode()
                self.send_response(status)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(raw)))
                self.end_headers()
                self.wfile.write(raw)

            def do_POST(self):
                seen["keys"].append(self.headers.get("X-API-Key"))
                length = int(self.headers.get("Content-Length") or 0)
                body = self.rfile.read(length) if length else b""
                if self.path == "/uploads":
                    seen["uploads"].append(body)
                    self._json(201, {"upload_id": "u1", "files": []})
                elif self.path == "/tasks":
                    seen["tasks"].append(json.loads(body.decode()))
                    self._json(201, {"id": "t1"})
                elif self.path == "/tasks/t1/cancel":
                    self._json(200, {})
                else:
                    self._json(404, {"error": "not found"})

            def do_GET(self):
                if self.path == "/tasks/t1":
                    polls["t1"] += 1
                    status = "queued" if polls["t1"] == 1 else "completed"
                    self._json(200, {"id": "t1", "status": status,
                                     "result_text": "Total: 85.15" if status == "completed" else "",
                                     "error_text": ""})
                elif self.path == "/tasks/t1/cost":
                    self._json(200, {"task_id": "t1", "usd_micros": 2500000})
                elif self.path == "/tasks/t1/trace":
                    self._json(200, {"entries": [{"id": "e1", "kind": "agent",
                                                  "text": "hello"}]})
                else:
                    self._json(404, {"error": "not found"})

        server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        try:
            client = HttpTaskClient(f"http://127.0.0.1:{server.server_address[1]}",
                                    "test-key")
            done = run_task(client, project_id="p1", prompt="do it",
                            files={"a.txt": b"hi"}, poll_seconds=0.01,
                            timeout_seconds=10)
        finally:
            server.shutdown()
            server.server_close()
        self.assertEqual(done.state.status, "completed")
        self.assertEqual(done.state.result_text, "Total: 85.15")
        self.assertEqual(done.cost_usd, 2.5)
        self.assertIn("hello", done.trace_text)
        self.assertTrue(all(k == "test-key" for k in seen["keys"]))
        self.assertEqual(seen["tasks"][0]["project_id"], "p1")
        fields, files = parse_multipart(seen["uploads"][0],
                                        "multipart/form-data; boundary=" +
                                        seen["uploads"][0].split(b"\r\n")[0][2:].decode())
        self.assertEqual(dict(files), {"a.txt": b"hi"})


class TestHttpClientRetry(unittest.TestCase):
    """瞬断重试：掉连接与 503 各重试一次后成功；401 只请求一次。"""

    def _serve(self, handler_cls):
        server = ThreadingHTTPServer(("127.0.0.1", 0), handler_cls)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        self.addCleanup(server.shutdown)
        self.addCleanup(server.server_close)
        return HttpTaskClient(
            f"http://127.0.0.1:{server.server_address[1]}", "test-key",
            timeout=5)

    def test_drop_then_success(self):
        counts = {"n": 0}

        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass

            def do_GET(self):
                counts["n"] += 1
                if counts["n"] == 1:
                    self.close_connection = True
                    return
                raw = json.dumps({"ok": True}).encode()
                self.send_response(200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(raw)))
                self.end_headers()
                self.wfile.write(raw)

        client = self._serve(Handler)
        self.assertEqual(client._request("GET", "/flaky"), {"ok": True})
        self.assertEqual(counts["n"], 2)

    def test_503_then_success(self):
        counts = {"n": 0}

        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass

            def do_GET(self):
                counts["n"] += 1
                raw = json.dumps({"ok": counts["n"] > 1}).encode()
                self.send_response(503 if counts["n"] == 1 else 200)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(raw)))
                self.end_headers()
                self.wfile.write(raw)

        client = self._serve(Handler)
        self.assertEqual(client._request("GET", "/flaky"), {"ok": True})
        self.assertEqual(counts["n"], 2)

    def test_401_not_retried(self):
        counts = {"n": 0}

        class Handler(BaseHTTPRequestHandler):
            def log_message(self, *args):
                pass

            def do_GET(self):
                counts["n"] += 1
                raw = json.dumps({"error": "unauthorized"}).encode()
                self.send_response(401)
                self.send_header("Content-Type", "application/json")
                self.send_header("Content-Length", str(len(raw)))
                self.end_headers()
                self.wfile.write(raw)

        client = self._serve(Handler)
        with self.assertRaises(RuntimeError):
            client._request("GET", "/flaky")
        self.assertEqual(counts["n"], 1)


class TestSweep(unittest.TestCase):
    def _spec(self, tag, **over):
        spec = {
            "run_tag": tag, "task_id": "ledger-1",
            "prompt": "Sum every signed amount in ledger.csv. Return Total: with two decimals.",
            "base_harness": dict(BASE_TREE),
            "source_files": {"ledger.csv": LEDGER},
            "k": 1, "mode": "parallel", "domain": "accounting",
        }
        spec.update(over)
        return spec

    def test_sweep_isolates_failures_and_shares_library(self):
        client = FakeTaskClient(fail_tags=("bad1",))
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        report = run_sweep(
            client, project_id="p1",
            runs=[self._spec("good1"), self._spec("bad1"),
                  self._spec("badk", k=99), self._spec("good2")],
            work_root=root / "sweep", library_root=root / "library")
        by_tag = {e["run_tag"]: e for e in report["runs"]}
        self.assertTrue(by_tag["good1"]["ok"])
        self.assertEqual(by_tag["good1"]["chosen_label"], "agg1")
        # bad1：全部任务注入失败 → 轮次 ok=False，但不中断后续轮。
        self.assertFalse(by_tag["bad1"]["ok"])
        # badk：k 越界在方法层抛错 → 记入 error，同样不中断。
        self.assertFalse(by_tag["badk"]["ok"])
        self.assertIn("ValueError", by_tag["badk"]["error"])
        self.assertTrue(by_tag["good2"]["ok"])
        self.assertEqual(report["totals"]["n_runs"], 4)
        self.assertEqual(report["totals"]["n_ok"], 2)
        # 共享库：两轮好结果都入库，同一 INDEX 可见。
        self.assertIsNotNone(by_tag["good1"]["library_entry"])
        self.assertIsNotNone(by_tag["good2"]["library_entry"])
        self.assertTrue((root / "library" / "accounting" / "INDEX.md").is_file())

    def test_sweep_requires_run_tag(self):
        client = FakeTaskClient()
        tmp = tempfile.TemporaryDirectory()
        self.addCleanup(tmp.cleanup)
        root = Path(tmp.name)
        spec = self._spec("good1")
        del spec["run_tag"]
        report = run_sweep(client, project_id="p1", runs=[spec],
                           work_root=root / "sweep")
        self.assertFalse(report["runs"][0]["ok"])
        self.assertIn("run_tag", report["runs"][0]["error"])


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