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 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 = [ "x\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```", "x\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 = [ "e\n\n```bash\nsed -i 's/a/b/' src/foo.py\n```", "s\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```", "e2\n\n```bash\nsed -i 's/b/c/' src/foo.py\n```", "s2\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 = "s\n\n```bash\necho SUBMIT_TASK_AAAA && cat patch.txt\n```" second = "s\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 = "r\n\n```bash\ncat src/foo.py\n```" edit = "e\n\n```bash\nsed -i 's/a/b/' src/foo.py\n```" edit2 = "e2\n\n```bash\nsed -i 's/b/c/' src/foo.py\n```" submit = "s\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 = [ "r\n\n```bash\ncat src/foo.py\n```", "x\n\n```bash\necho ALBEDO_TASK_DONE_SUBMIT_NOW\n```", "x\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 = "s\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 = "s\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 = [ "e\n\n```bash\nsed -i '/def test_wheel_tag_expand/,/^def / { /mac_ver/a\\ x' tests/test_builder.py\n```", "s\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```", "g\n\n```bash\ngrep -n maxsize tests/test_builder.py\n```", "c\n\n```bash\ncat -n tests/test_builder.py | grep test_wheel_tag_expand\n```", "s\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```", "e\n\n```bash\nsed -i '/def test_wheel_tag_expand_11/,/^def / { /mac_ver/a\\ x' tests/test_builder.py\n```", "s\n\n```bash\necho COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT\n```", "g\n\n```bash\ngrep -rn buildver src/scikit_build_core --include='*.py'\n```", "s\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\n" chosen = "\nbody\n\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