code/tests/test_local_train_chain.py

import json
import random
from pathlib import Path
from types import SimpleNamespace

from albedo_eval_service.shared.observation_format import MAX_CONSECUTIVE_BAD_TURNS
from local_eval.constants import MAX_NEW_TOKENS
from local_train.chain_eval import pick_free_gpus
from local_train.chain_heuristics import (
    ChainState,
    empty_submit_count_from_texts,
    evaluate_chain,
    mid_roll_fatal,
)
from albedo_eval_service.shared.observation_format import first_bash_block
from local_train.chain_gold import compile_session, micro_ok, path_aliases, session_passes
from local_train.chain_pack import _mix_stage_a, followup_edit_command, infer_micro
from local_train.reward import score_texts
from sanity_service.chain import empty_submit_count


def _state(**kwargs) -> ChainState:
    defaults = dict(
        sample_id="s1",
        submit_clause="echo SUBMIT_TASK_AAAA && cat patch.txt",
        submit_marker="SUBMIT_TASK_AAAA",
        micro={"file": "src/foo.py", "function": "bar", "request": "edit bar", "message": ""},
    )
    defaults.update(kwargs)
    return ChainState(**defaults)


def test_infer_micro_reads_path_and_function():
    micro = infer_micro("please look at src/util/foo.py\ndef bar():\n    pass")
    assert micro["file"] == "src/util/foo.py"
    assert micro["function"] == "bar"
    assert "foo.py" in followup_edit_command(micro["file"], micro["function"])


def test_infer_micro_skips_english_function_to():
    hay = (
        "Function: debounce<F extends (...args: any) => any>(func: F, options: {})\n"
        "Location: src/debounce.ts\n"
        "the function to debounce; its parameter and return types are preserved.\n"
    )
    micro = infer_micro(hay)
    assert micro["file"] == "src/debounce.ts"
    assert micro["function"] == "debounce"


def test_infer_micro_pairs_bold_symbol_with_matching_file():
    hay = (
        "Throws **SlackError** when parsing fails.\n"
        "Returns a parsed **Content** object.\n"
        "Edit src/content.js and src/errors.js.\n"
    )
    micro = infer_micro(hay)
    assert micro["file"] == "src/content.js"
    assert micro["function"] == "Content"


def test_infer_micro_prefers_cmakelists_on_cmake_pr():
    hay = (
        "# Build fails with GCC on macOS due to forced `-stdlib=libc++` flag\n"
        "CMakeLists.txt unconditionally adds the flag when APPLE.\n"
        "set(FlatBuffers_Library_SRCS\n"
        "  include/flatbuffers/base.h\n"
        "  include/flatbuffers/flatbuffers.h\n"
        ")\n"
        "Use CMAKE_CXX_FLAGS only for Clang.\n"
    )
    micro = infer_micro(hay)
    assert micro["file"].endswith("CMakeLists.txt")
    assert micro["function"] != "the helper"


def test_two_empty_submits_is_the_live_fault():
    texts = [
        "<think>x</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```",
        "<think>x</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```",
    ]
    assert empty_submit_count_from_texts(texts, "SUBMIT_TASK_AAAA") >= 2
    ns = SimpleNamespace(
        turns=[{"role": "assistant", "content": t, "score_target": True} for t in texts]
    )
    assert empty_submit_count(ns, "SUBMIT_TASK_AAAA") >= 2


def test_edit_between_submits_is_not_empty():
    texts = [
        "<think>e</think>\n\n```bash\nsed -i 's/a/b/' src/foo.py\n```",
        "<think>s</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```",
        "<think>e2</think>\n\n```bash\nsed -i 's/b/c/' src/foo.py\n```",
        "<think>s2</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```",
    ]
    assert empty_submit_count_from_texts(texts, "SUBMIT_TASK_AAAA") == 0


def test_evaluate_chain_flags_empty_double_submit():
    first = "<think>s</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```"
    second = "<think>s</think>\n\n```bash\necho SUBMIT_TASK_AAAA\n```"
    state = _state(
        turns=[
            {"role": "user", "content": "Please handle the missed call site.", "segment": "micro"},
            {"role": "assistant", "content": first, "score_target": True, "segment": "micro"},
            {"role": "user", "content": "Please handle the missed call site.", "injected": True},
            {"role": "assistant", "content": second, "score_target": True, "segment": "micro"},
        ],
        submits=[{"post_nudge": False, "has_edit": False}, {"post_nudge": False, "has_edit": False}],
        micro=None,
    )
    evaluate_chain([state], 8)
    assert state.heuristic_reason == "chain: repeated submissions without doing any work"


def test_evaluate_chain_passes_edit_then_submit_then_edit_then_submit():
    read = "<think>r</think>\n\n```bash\ncat src/foo.py\n```"
    edit = "<think>e</think>\n\n```bash\nsed -i 's/a/b/' src/foo.py\n```"
    edit2 = "<think>e2</think>\n\n```bash\nsed -i 's/b/c/' src/foo.py\n```"
    submit = "<think>s</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```"
    state = _state(
        turns=[
            {"role": "user", "content": "handle src/foo.py bar; echo SUBMIT_TASK_AAAA", "segment": "micro"},
            {"role": "assistant", "content": read, "score_target": True, "segment": "micro"},
            {"role": "user", "content": "ok"},
            {"role": "assistant", "content": edit, "score_target": True, "segment": "micro"},
            {"role": "user", "content": "ok"},
            {"role": "assistant", "content": submit, "score_target": True, "segment": "micro"},
            {"role": "user", "content": "not quite"},
            {"role": "assistant", "content": edit2, "score_target": True, "segment": "micro"},
            {"role": "user", "content": "ok"},
            {"role": "assistant", "content": submit, "score_target": True, "segment": "micro"},
        ],
        submits=[{"post_nudge": False, "has_edit": True}, {"post_nudge": False, "has_edit": True}],
    )
    evaluate_chain([state], 8)
    assert state.heuristic_reason == ""


def test_reward_zero_on_empty_double_submit():
    texts = [
        "<think>r</think>\n\n```bash\ncat src/foo.py\n```",
        "<think>x</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW\n```",
        "<think>x</think>\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW\n```",
    ]
    got = score_texts(texts, submit_command="echo ALBEDO_TASK_DONE_SUBMIT_NOW")
    assert got.fatal and got.reward == 0.0
    assert "empty_double_submit" in got.reasons


def test_mix_stage_a_prefers_at_edit_and_post_edit_submit(tmp_path: Path):
    rows = [
        {"sample_id": "e1", "kind": "edit", "phase": "at_edit", "prompt": "p", "completion": "c"},
        {"sample_id": "e2", "kind": "edit", "phase": "explore", "prompt": "p", "completion": "c"},
        {"sample_id": "s1", "kind": "submit", "phase": "post_edit", "prompt": "p", "completion": "c"},
        {"sample_id": "s2", "kind": "submit", "phase": "cold", "prompt": "p", "completion": "c"},
        {"sample_id": "x1", "kind": "explore", "phase": "explore", "prompt": "p", "completion": "c"},
    ]
    pack = tmp_path / "mix.jsonl"
    pack.write_text("".join(json.dumps(row) + "\n" for row in rows))
    mixed = _mix_stage_a(pack, random.Random(0), edit_n=1, submit_n=1)
    assert {row.sample_id for row in mixed} == {"e1", "s1"}


def test_micro_ok_drops_weak_targets():
    assert not micro_ok({"file": "content.js", "function": "bar"})
    assert not micro_ok({"file": "src/foo.py", "function": "the helper"})
    assert not micro_ok({"file": "src/foo.py", "function": "to"})
    assert micro_ok({"file": "src/util/foo.py", "function": "bar"})
    assert micro_ok({"file": "CMakeLists.txt", "function": "GCC"})


def test_mid_roll_fatal_flags_loop_without_waiting_for_horizon():
    submit = "<think>s</think>\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```"
    state = _state(
        turns=[
            {"role": "assistant", "content": submit, "score_target": True, "segment": "micro"}
            for _ in range(5)
        ],
        submits=[{"post_nudge": False, "has_edit": False} for _ in range(5)],
    )
    assert "same command repeated" in mid_roll_fatal(state)


def test_pick_free_gpus_skips_occupied_cards(monkeypatch):
    monkeypatch.setattr(
        "local_train.chain_eval.subprocess.check_output",
        lambda *args, **kwargs: "0, 20000\n1, 120000\n2, 130000\n3, 140000\n4, 5000\n",
    )
    assert pick_free_gpus(3, min_free_gib=100) == ["1", "2", "3"]


def test_pick_free_gpus_errors_when_not_enough(monkeypatch):
    monkeypatch.setattr(
        "local_train.chain_eval.subprocess.check_output",
        lambda *args, **kwargs: "0, 20000\n1, 21000\n",
    )
    try:
        pick_free_gpus(4, min_free_gib=100)
    except RuntimeError as exc:
        assert "need 4 GPUs" in str(exc)
    else:
        raise AssertionError("expected RuntimeError")


def test_kill_process_tree_reaps_child():
    import subprocess as sp

    from albedo_eval_service.remote.generation import _kill_process_tree

    proc = sp.Popen(["sleep", "60"])
    _kill_process_tree(proc.pid)
    proc.wait(timeout=3)
    assert proc.poll() is not None


def test_compiled_session_passes_evaluate_chain():
    micro = {
        "file": "src/util/foo.py",
        "function": "bar",
        "request": "edit bar in foo.py",
        "message": "",
    }
    clause = "echo SUBMIT_TASK_AAAA && cat patch.txt"
    completions = compile_session(micro, clause)
    assert completions is not None
    assert len(completions) >= 20
    assert session_passes(micro, clause, "SUBMIT_TASK_AAAA", completions) == ""
    cmds = [first_bash_block(text) or "" for text in completions]
    assert all("sed" not in cmd or "&&" not in cmd for cmd in cmds)
    for prev, cur in zip(cmds, cmds[1:]):
        assert not ("SUBMIT_TASK_AAAA" in prev and "SUBMIT_TASK_AAAA" in cur)


def test_live_token_cap_is_4k():
    assert MAX_NEW_TOKENS == 4096
    assert MAX_CONSECUTIVE_BAD_TURNS == 3


def test_is_submit_turn_accepts_marker_without_exact_tail():
    from local_train.chain_heuristics import is_submit_turn

    text = "<think>s</think>\n\n```bash\necho SUBMIT_TASK_AAAA\n```"
    assert is_submit_turn(text, "echo SUBMIT_TASK_AAAA && cat patch.txt", "SUBMIT_TASK_AAAA")


def test_path_aliases_include_testbed_prefix():
    aliases = path_aliases("CMakeLists.txt")
    assert "CMakeLists.txt" in aliases
    assert "testbed/CMakeLists.txt" in aliases
    assert "/testbed/CMakeLists.txt" in aliases
    completions = compile_session(
        {"file": "CMakeLists.txt", "function": "GCC", "request": "edit", "message": ""},
        "echo FINALIZE_AND_SUBMIT_TASK_OUTPUT",
        path="testbed/CMakeLists.txt",
    )
    assert completions is not None
    assert any("cat testbed/CMakeLists.txt" in text for text in completions)
    assert any("sed -i" in text and "testbed/CMakeLists.txt" in text for text in completions)


def test_live_nosa_trace_is_empty_double_submit():
    """Exact live v11/nosa pattern: sed, bare echo, grep/cat, bare echo, sed, echo, grep, echo."""
    marker = "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT"
    texts = [
        "<think>e</think>\n\n```bash\nsed -i '/def test_wheel_tag_expand/,/^def / { /mac_ver/a\\    x' tests/test_builder.py\n```",
        "<think>s</think>\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```",
        "<think>g</think>\n\n```bash\ngrep -n maxsize tests/test_builder.py\n```",
        "<think>c</think>\n\n```bash\ncat -n tests/test_builder.py | grep test_wheel_tag_expand\n```",
        "<think>s</think>\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```",
        "<think>e</think>\n\n```bash\nsed -i '/def test_wheel_tag_expand_11/,/^def / { /mac_ver/a\\    x' tests/test_builder.py\n```",
        "<think>s</think>\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```",
        "<think>g</think>\n\n```bash\ngrep -rn buildver src/scikit_build_core --include='*.py'\n```",
        "<think>s</think>\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```",
    ]
    assert empty_submit_count_from_texts(texts, marker) >= 2


def test_reject_edit_thought_forbids_submit_now():
    from local_train.chain_gold import reject_edit_thought

    thought = reject_edit_thought("src/foo.py", "bar", "echo SUBMIT_TASK_AAAA && cat patch.txt")
    assert "not submit-now" in thought
    assert "sed -i" in thought
    assert "miner fault" in thought


def test_path_only_session_passes_on_live_micros():
    clause = "echo FINALIZE_AND_SUBMIT_TASK_OUTPUT"
    for micro in (
        {"file": "src/content.js", "function": "Content", "request": "edit Content", "message": ""},
        {"file": "CMakeLists.txt", "function": "GCC", "request": "edit CMakeLists", "message": ""},
        {"file": "src/debounce.ts", "function": "debounce", "request": "edit debounce", "message": ""},
    ):
        completions = compile_session(micro, clause)
        assert completions is not None
        assert session_passes(micro, clause, "FINALIZE_AND_SUBMIT_TASK_OUTPUT", completions) == ""
        assert all("sed -i" not in text or clause not in text for text in completions)
        cmds = [first_bash_block(text) or "" for text in completions]
        assert len(cmds) >= 20
        for prev, cur in zip(cmds, cmds[1:]):
            assert not ("FINALIZE_AND_SUBMIT_TASK_OUTPUT" in prev and "FINALIZE_AND_SUBMIT_TASK_OUTPUT" in cur)
        assert sum("sed -i" in cmd for cmd in cmds) >= 8
        assert sum(clause.split()[0] in cmd for cmd in cmds) == 2


def test_compiled_session_has_exactly_two_submits():
    micro = {
        "file": "src/util/foo.py",
        "function": "bar",
        "request": "edit bar",
        "message": "",
    }
    clause = "echo SUBMIT_TASK_AAAA && cat patch.txt"
    completions = compile_session(micro, clause)
    assert completions is not None
    cmds = [first_bash_block(text) or "" for text in completions]
    assert cmds[2] == clause
    assert cmds[4] == clause
    assert all("SUBMIT_TASK_AAAA" not in cmd for cmd in cmds[5:])
    assert all("sed -i" in cmd for i, cmd in enumerate(cmds) if i not in {0, 2, 4})


def test_dpo_pairs_prefer_sed_over_live_ban():
    from local_train.dpo_pack import CLAUSES, iter_prefixes, strip_open_think

    prefixes = iter_prefixes()
    assert {p.kind for p in prefixes} >= {"nosa", "mixed", "messy", "recat", "after_submit"}
    assert len({p.sample_id for p in prefixes}) == len(prefixes)
    assert any(p.file.startswith("tests/") for p in prefixes)
    assert any("cmake.org" in p.file for p in prefixes)
    nosa = next(p for p in prefixes if p.kind == "nosa")
    cmds = []
    for msg in nosa.messages:
        if msg["role"] == "assistant":
            cmds.append(first_bash_block(msg["content"]) or "")
    assert any(c.startswith("sed -i") for c in cmds)
    assert any("grep" in c for c in cmds)
    assert any(nosa.submit_marker in c for c in cmds)
    last_user = nosa.messages[-1]["content"]
    assert "Submit the same way when done" in last_user
    assert nosa.submit_command in last_user
    mixed = next(p for p in prefixes if p.kind == "mixed")
    last_asst = first_bash_block(mixed.messages[-2]["content"]) or ""
    assert "sed" in last_asst and "&&" in last_asst and mixed.submit_marker in last_asst
    for prefix in prefixes:
        assert "sed -i" in prefix.chosen_bash
        assert prefix.submit_marker not in prefix.chosen_bash
        assert any(prefix.submit_marker in bash for bash in prefix.rejected_bashes)
        assert prefix.chosen_bash not in prefix.rejected_bashes
    prompt = "assistant\n<think>\n"
    chosen = "<think>\nbody\n</think>\n\n```bash\nsed -i '1a\\x' a.py\n```\n"
    assert strip_open_think(prompt, chosen).startswith("body")
    assert all(marker for _, marker in CLAUSES)


def test_late_horizon_trap_is_sed_after_a_submit():
    from types import SimpleNamespace

    from local_eval.constants import TOKENIZER_DIR
    from local_train.chain_gold import _gold_rows

    micro = {
        "file": "src/util/foo.py",
        "function": "bar",
        "request": "edit bar",
        "message": "",
    }
    clause = "echo SUBMIT_TASK_AAAA && cat patch.txt"
    completions = compile_session(micro, clause)
    assert completions is not None
    rows = _gold_rows(
        sample_id="late-trap-test",
        prefix=[],
        micro=micro,
        completions=completions,
        clause=clause,
        marker="SUBMIT_TASK_AAAA",
        rewrite_mode="test",
        item=SimpleNamespace(source="test", gold="src/util/foo.py", family="chain", language="", repo=""),
        salt="late-test",
        tokenizer=str(TOKENIZER_DIR),
    )
    late = [row for row in rows if row.kind == "chain_late_trap"]
    assert late
    for row in late:
        bash = first_bash_block(row.completion) or ""
        assert "sed -i" in bash
        assert "SUBMIT_TASK_AAAA" not in bash
        assert "not submit-now" in row.completion
        assert "SUBMIT_TASK_AAAA" in row.prompt
        assert "chain-late" in bash