tests/test_edit_loop.py

"""act → verify → correct: the half of the loop that was missing.

Every other tool in this engine reads. Turing could describe an edit in
perfect detail and never make one — the scientist still opened the editor and
typed it themselves. That gap is the whole difference between an agent and a
very good assistant.

This is also the first thing in the codebase that CHANGES a construct, so the
tests here are weighted toward refusal rather than capability. The dangerous
outcome is not "the edit failed"; it is "the edit succeeded, somewhere else,
quietly", and someone orders that DNA.

Four properties, each of which fails silently if broken:

  1. An edit that does not match the construct is REFUSED, and the refusal
     says what is actually there.
  2. Residue numbering is only used on an actual reading frame.
  3. The edited construct becomes the run's target — otherwise the verify
     step re-checks the old sequence and reports a clean bill of health for a
     change that never landed.
  4. The bases never travel through the model.
"""
import json
import time

import pytest

from dee.core import agent_tools as t
from dee.core import edits
from dee.core import llm as _llm
from dee.core import orchestrator as orch


# A real, minimal CDS: M R A K stop.  Residue 2 is R (CGT).
CDS = "ATG" + "CGT" + "GCA" + "AAA" + "TAA"


# --------------------------------------------------------------------------- #
# the engine
# --------------------------------------------------------------------------- #
def test_a_substitution_lands_on_the_right_codon():
    out = edits.apply_edits(CDS, "R2H")
    assert out["ok"] is True
    assert edits.translate(out["sequence"]) == "MHAK*"
    a = out["applied"][0]
    assert a["codon_before"] == "CGT"
    assert (a["dna_from"], a["dna_to"]) == (4, 6)


def test_the_codon_that_was_installed_is_reported():
    """H is CAT or CAC and the choice affects expression. Leaving it implicit
    would make the tool look more deterministic than it is."""
    out = edits.apply_edits(CDS, "R2H", host="e_coli")
    note = out["applied"][0]["codon_note"]
    assert out["applied"][0]["codon_after"] in ("CAT", "CAC")
    assert out["applied"][0]["codon_after"] in note
    assert "e_coli" in note


def test_a_wild_type_mismatch_is_refused_and_says_what_is_there():
    """THE test. Applying K2H to a construct whose residue 2 is R would
    mutate the wrong residue and look like a success. The refusal has to name
    the actual residue, or the agent cannot tell a numbering problem from a
    typo."""
    out = edits.apply_edits(CDS, "K2H")
    assert out["ok"] is False
    assert out["kind"] == "mismatch"
    assert "is R" in out["error"] and "not K" in out["error"]


def test_residue_numbering_is_refused_on_a_non_coding_sequence():
    """A residue number quoted against something that is not a reading frame
    is a number about nothing. Measuring from base 0 and hoping is how an
    edit lands 200 bp from where it was meant to."""
    out = edits.apply_edits("ACGTACGTACGTACGTAC", "R2H")
    assert out["ok"] is False
    assert out["kind"] == "not_coding"


def test_nothing_is_applied_unless_every_edit_validates():
    """A half-applied set is a construct nobody asked for, and the agent has
    no way to tell which half it got."""
    out = edits.apply_edits(CDS, ["R2H", "K99W"])
    assert out["ok"] is False
    # and the good one did not sneak through
    assert "sequence" not in out


def test_two_edits_at_one_position_are_refused():
    out = edits.apply_edits(CDS, ["R2H", "R2W"])
    assert out["ok"] is False
    assert "position 2" in out["error"]


def test_the_level_is_never_inferred():
    """'A123G' is a valid edit in both alphabets. Guessing which one was meant
    would silently edit the wrong thing on exactly the constructs where it
    matters most, so `level` is required and the parser respects it."""
    assert edits.parse_edit("A12G", "dna") == {"from": "A", "pos": 12, "to": "G"}
    assert edits.parse_edit("A12G", "protein") == {"from": "A", "pos": 12, "to": "G"}
    # ...but a residue that is not a base is rejected at the dna level
    assert edits.parse_edit("R12H", "dna") is None


def test_a_base_level_edit_works_on_anything():
    out = edits.apply_edits("ACGTACGT", "G3T", level="dna")
    assert out["ok"] is True
    assert out["sequence"] == "ACTTACGT"
    assert out["length_changed"] == 0


def test_a_base_level_mismatch_is_refused_too():
    out = edits.apply_edits("ACGTACGT", "A3T", level="dna")
    assert out["ok"] is False and out["kind"] == "mismatch"
    assert "is G" in out["error"]


def test_a_wholesale_replacement_is_refused_as_an_edit():
    out = edits.apply_edits(CDS, [f"R{i}H" for i in range(1, 40)])
    assert out["ok"] is False
    assert "cap is" in out["error"]


def test_is_coding_is_strict():
    assert edits.is_coding(CDS) is True
    assert edits.is_coding(CDS[:-1]) is False        # not whole codons
    assert edits.is_coding("CGT" * 4 + "TAA") is False   # no ATG
    assert edits.is_coding("ATG" + "TAA" + "GCA" + "TAA") is False  # internal stop


# --------------------------------------------------------------------------- #
# the tool
# --------------------------------------------------------------------------- #
def test_the_tool_returns_the_construct_to_the_browser_only():
    """_strip_ui removes `_ui` before the model sees the result. The bases
    belong to the editor; in the model's context they are pure cost, and for
    a construct over ~8 kB they would be truncated anyway."""
    out = t.execute_tool("edit_sequence",
                         {"sequence": CDS, "edits": ["R2H"], "level": "protein"},
                         auth_anonymous=False, user_id="u" * 36)
    assert out["ok"] is True
    assert out["_ui"]["sequence"]
    stripped = orch._strip_ui(out)
    assert "_ui" not in stripped
    assert not any(isinstance(v, str) and len(v) > 40 and set(v) <= set("ACGT")
                   for v in stripped.values()), "bases leaked into the model result"


def test_the_result_tells_the_agent_to_verify():
    """Carried in the RESULT, not only in the system prompt, so it survives
    compaction and arrives at the moment it is relevant."""
    out = t.execute_tool("edit_sequence",
                         {"sequence": CDS, "edits": ["R2H"], "level": "protein"},
                         auth_anonymous=False, user_id="u" * 36)
    assert "map_plasmid" in out["next"]
    assert "verify" in out["next"].lower()


def test_the_edited_spans_are_reported_for_the_editor():
    """So the editor can select what changed. Repainting a 20 kb construct
    that differs by one base and leaving the viewport at base 1 asks the
    scientist to take the edit on trust."""
    out = t.execute_tool("edit_sequence",
                         {"sequence": CDS, "edits": ["R2H"], "level": "protein"},
                         auth_anonymous=False, user_id="u" * 36)
    spans = out["_ui"]["edited_spans"]
    assert spans == [{"from": 4, "to": 6, "label": "R2H"}]


def test_the_tool_has_no_sequence_parameter():
    """The sequence is injected by the orchestrator from the bound target.
    If the schema advertised one, the model would fill it — from a context
    where a large construct has already been truncated to 8 kB."""
    spec = next(s for s in orch.TOOL_SPECS
                if s["function"]["name"] == "edit_sequence")
    props = spec["function"]["parameters"]["properties"]
    assert "sequence" not in props
    assert set(spec["function"]["parameters"]["required"]) == {"edits", "level"}


def test_the_tool_is_gated_and_target_bound():
    assert t._TOOLS["edit_sequence"]["requires_confirm"] is True
    assert t.needs_target("edit_sequence") is True
    assert t.needs_target("map_plasmid") is False


# --------------------------------------------------------------------------- #
# the loop, through the orchestrator
# --------------------------------------------------------------------------- #
def _msg_tool(name, args, call_id="c1"):
    return {"content": None, "tool_calls": [
        {"id": call_id, "function": {"name": name, "arguments": json.dumps(args)}}]}


def _msg_text(text):
    return {"content": text, "tool_calls": []}


def _script(monkeypatch, messages):
    seq = list(messages)

    def fake(config, body):
        return (seq.pop(0) if seq else _msg_text("done")), 0.001, 10, 0
    monkeypatch.setattr(_llm, "call", fake)


def _rest(run, timeout=5.0):
    deadline = time.time() + timeout
    while time.time() < deadline:
        if run.status in ("awaiting_input", "done", "error", "stopped"):
            return run.status
        time.sleep(0.01)
    raise AssertionError(f"never settled ({run.status})")


@pytest.fixture(autouse=True)
def _configured(monkeypatch):
    monkeypatch.setenv("OPENROUTER_API_KEY", "test-key")
    monkeypatch.setenv("AGENT_MAX_COST_USD", "0.5")


def test_the_sequence_is_injected_not_passed(monkeypatch):
    """The model calls edit_sequence with no bases; the orchestrator supplies
    them from run.target. This is what makes a 19,895 bp construct editable
    at all."""
    seen = {}

    def exec_tool(name, args, anon, **kw):
        seen.update(args)
        return {"ok": True, "_ui": {"sequence": "ATGCATGCAAAATAA"},
                "edits_applied": ["R2H"], "length": 15}
    monkeypatch.setattr(orch._tools, "execute_tool", exec_tool)
    _script(monkeypatch, [_msg_tool("edit_sequence",
                                    {"edits": ["R2H"], "level": "protein"}),
                          _msg_text("done")])
    run = orch.create_run(owner="u1", anonymous=False)
    run.target = {"label": "TP53", "gene_symbol": "TP53", "sequence": CDS,
                  "length": len(CDS), "uniprot": "P04637"}
    orch.start(run, "install R2H")
    assert _rest(run) == "awaiting_input", "the gate should have stopped it"
    orch.reply(run, "yes")
    _rest(run)
    assert seen.get("sequence") == CDS, "the construct was not injected"


def test_editing_with_nothing_loaded_fails_cleanly(monkeypatch):
    """And answers its tool call, so history stays valid."""
    monkeypatch.setattr(orch._tools, "execute_tool",
                        lambda *a, **k: pytest.fail("must not reach the tool"))
    _script(monkeypatch, [_msg_tool("edit_sequence",
                                    {"edits": ["R2H"], "level": "protein"},
                                    call_id="e1"),
                          _msg_text("Nothing loaded.")])
    run = orch.create_run(owner="u1", anonymous=False)
    orch.start(run, "install R2H")
    _rest(run)
    reply = next(m for m in run.history
                 if m.get("role") == "tool" and m["tool_call_id"] == "e1")
    assert json.loads(reply["content"])["kind"] == "no_target"


def test_the_edited_construct_becomes_the_target(monkeypatch):
    """Without this the verify step re-checks the sequence as it was BEFORE
    the edit and reports it healthy — a clean bill of health for a change
    that never entered the record."""
    edited = "ATG" + "CAT" + "GCA" + "AAA" + "TAA"

    def exec_tool(name, args, anon, **kw):
        return {"ok": True, "_ui": {"sequence": edited},
                "edits_applied": ["R2H"], "length": len(edited)}
    monkeypatch.setattr(orch._tools, "execute_tool", exec_tool)
    _script(monkeypatch, [_msg_tool("edit_sequence",
                                    {"edits": ["R2H"], "level": "protein"}),
                          _msg_text("done")])
    run = orch.create_run(owner="u1", anonymous=False)
    run.target = {"label": "TP53", "gene_symbol": "TP53", "sequence": CDS,
                  "length": len(CDS), "uniprot": "P04637"}
    orch.start(run, "install R2H")
    _rest(run); orch.reply(run, "yes"); _rest(run)
    assert run.target["sequence"] == edited
    assert run.target["edited"] is True
    assert "R2H" in run.target["label"]


def test_the_wild_type_accession_is_dropped_after_an_edit(monkeypatch):
    """P04637 describes wild-type TP53. Keeping it on a mutant would have the
    structure viewer confidently show an unmutated model for an edited
    construct — the TEM-1 incident's failure mode, one edit later."""
    def exec_tool(name, args, anon, **kw):
        return {"ok": True, "_ui": {"sequence": "ATGCATGCAAAATAA"},
                "edits_applied": ["R2H"], "length": 15}
    monkeypatch.setattr(orch._tools, "execute_tool", exec_tool)
    _script(monkeypatch, [_msg_tool("edit_sequence",
                                    {"edits": ["R2H"], "level": "protein"}),
                          _msg_text("done")])
    run = orch.create_run(owner="u1", anonymous=False)
    run.target = {"label": "TP53", "gene_symbol": "TP53", "sequence": CDS,
                  "length": len(CDS), "uniprot": "P04637"}
    orch.start(run, "install R2H")
    _rest(run); orch.reply(run, "yes"); _rest(run)
    assert not run.target.get("uniprot")


def test_the_lineage_does_not_compound_over_repeated_edits(monkeypatch):
    """Two edits must not produce 'TP53 + R2H + R2H + A3G'. The stem is kept
    and the labels replaced, so the name stays readable."""
    def exec_tool(name, args, anon, **kw):
        return {"ok": True, "_ui": {"sequence": "ATGCATGCAAAATAA"},
                "edits_applied": ["A3G"], "length": 15}
    monkeypatch.setattr(orch._tools, "execute_tool", exec_tool)
    _script(monkeypatch, [_msg_tool("edit_sequence",
                                    {"edits": ["A3G"], "level": "protein"}),
                          _msg_text("done")])
    run = orch.create_run(owner="u1", anonymous=False)
    run.target = {"label": "TP53 + R2H", "gene_symbol": "TP53",
                  "sequence": CDS, "length": len(CDS)}
    orch.start(run, "now A3G")
    _rest(run); orch.reply(run, "yes"); _rest(run)
    assert run.target["label"] == "TP53 + A3G"


def test_the_prompt_makes_verification_part_of_the_job():
    """A tool that changes something and a prompt that says "you're done" is
    how an unverified edit reaches a bench."""
    p = orch.build_system_prompt(False, {}, "")
    assert "edit_sequence" in p
    assert "NOT done when the tool returns" in p
    assert "map_plasmid" in p
    # and the refusal is framed as information, not an obstacle to route round
    assert "Never retry the same edit with a different guess" in p