LLeMbas/tests/test_agent_plan.py

"""Plans: the structure, the compatibility, and keeping one current.

The compatibility half is the one that matters most. Every plan row on disk is
`{title, steps}`, and `execute_plan` reads `steps` -- so the rule is that
`steps` is always written, and a version-1 row normalises into the new shape
rather than being migrated.
"""

from __future__ import annotations

import json as _json

import pytest

from lembas.db.models import KIND_AGENT, ROLE_ASSISTANT, Chat, Connection, Message, Model, User
from lembas.security.passwords import hash_password
from lembas.services import plans
from lembas.services import tools as tools_service
from lembas.services.agent import policy, session

# --- The shape ------------------------------------------------------------------
V1 = {"title": "Tidy the logs", "steps": ["Read the log", "Rotate it", "Restart"]}


def test_a_version_one_row_becomes_one_phase():
    """Every row that exists today. The old branch lives in `normalise` and
    nowhere else, so nothing downstream deals with two shapes."""
    plan = plans.normalise(V1)

    assert plan["version"] == 2
    assert plan["title"] == "Tidy the logs"
    assert len(plan["phases"]) == 1
    assert [t["text"] for t in plan["phases"][0]["tasks"]] == V1["steps"]
    assert plan["steps"] == V1["steps"], "what execute_plan reads, unchanged"


def test_steps_is_always_written_and_is_the_flattened_tasks():
    """The whole of the compatibility story: nothing downstream had to learn
    version 2."""
    plan = plans.build(
        title="Ship it",
        phases=[
            {"title": "Survey", "tasks": ["Read the config", "List the services"]},
            {"title": "Change", "tasks": ["Patch the unit file"]},
        ],
    )
    assert plan["steps"] == ["Read the config", "List the services", "Patch the unit file"]


def test_findings_and_objectives_survive():
    plan = plans.build(
        title="Ship it",
        findings=["The unit file is generated"],
        objectives=["It restarts cleanly"],
        steps=["Do the thing"],
    )
    assert plan["findings"][0]["text"] == "The unit file is generated"
    assert plan["objectives"][0]["status"] == "open"


def test_ids_are_ours_and_are_unique_across_phases():
    """Letting the model name them would mean validating names it made up, and
    a collision would silently re-tick a different task."""
    plan = plans.build(
        title="x",
        phases=[
            {"title": "One", "tasks": ["a", "b"], "id": "MINE"},
            {"title": "Two", "tasks": ["c"]},
        ],
    )
    ids = [t["id"] for phase in plan["phases"] for t in phase["tasks"]]
    assert ids == ["t1", "t2", "t3"]
    assert plan["phases"][0]["id"] == "p1"


def test_a_bare_string_where_a_list_was_expected_is_tolerated():
    """The same tolerance `_questions_in` shows. A small model sends something
    close to the schema, and refusing costs a whole round trip."""
    plan = plans.build(title="x", steps="just the one thing")
    assert plan["steps"] == ["just the one thing"]


def test_an_empty_plan_stays_empty():
    assert plans.normalise({}) == {}
    assert plans.normalise(None) == {}


# --- Updating --------------------------------------------------------------------
def _plan():
    return plans.build(
        title="Ship it",
        objectives=["It restarts cleanly"],
        phases=[
            {"title": "Survey", "tasks": ["Read the config", "List the services"]},
            {"title": "Change", "tasks": ["Patch the unit file"]},
        ],
    )


def test_marking_a_task_done_changes_it_and_says_so():
    plan, changed = plans.merge(_plan(), {"task_status": [{"id": "t1", "status": "done"}]})

    assert plan["phases"][0]["tasks"][0]["status"] == "done"
    assert "t1 is done" in changed


def test_an_unknown_id_changes_nothing():
    """And reports nothing, so the runner can tell the model to quote a real id
    rather than silently succeeding at nothing."""
    _plan_, changed = plans.merge(_plan(), {"task_status": [{"id": "t99", "status": "done"}]})
    assert changed == []


def test_a_finished_phase_collapses_and_the_next_becomes_active():
    """A phase's status follows from its tasks, so the two cannot disagree --
    a plan reading "phase 1: done" over four todo tasks is worse than either."""
    plan, _ = plans.merge(
        _plan(),
        {"task_status": [{"id": "t1", "status": "done"}, {"id": "t2", "status": "done"}]},
    )
    assert plan["phases"][0]["status"] == "done"
    assert plan["phases"][1]["status"] == "active"


def test_only_one_phase_is_ever_active():
    plan = plans.normalise(_plan())
    assert [p["status"] for p in plan["phases"]].count("active") <= 1


def test_a_task_added_mid_work_lands_in_the_phase_being_worked_on():
    plan, changed = plans.merge(_plan(), {"add_tasks": [{"text": "Back up the old one"}]})

    assert "Back up the old one" in [t["text"] for t in plan["phases"][0]["tasks"]]
    assert plan["steps"][-1] != "Back up the old one", "it goes in phase one, not at the end"
    assert changed


def test_a_new_finding_is_appended():
    plan, changed = plans.merge(_plan(), {"findings": ["The service is socket-activated"]})
    assert plan["findings"][-1]["text"] == "The service is socket-activated"
    assert changed


def test_updating_keeps_steps_in_step():
    plan, _ = plans.merge(_plan(), {"add_tasks": [{"text": "Back up the old one"}]})
    assert plan["steps"] == plans.flatten(plan)


# --- The block the model sees ------------------------------------------------------
def test_the_block_shows_the_ids_the_update_tool_takes():
    block = plans.render_block(_plan())
    assert "t1" in block and "o1" in block


def test_a_finished_phase_is_one_line_in_the_block():
    """Budgeted rather than dumped, exactly like the project listing."""
    plan, _ = plans.merge(
        _plan(),
        {"task_status": [{"id": "t1", "status": "done"}, {"id": "t2", "status": "done"}]},
    )
    block = plans.render_block(plan)

    assert "Survey (2 tasks, done)" in block
    assert "Read the config" not in block, "a finished phase collapses"
    assert "Patch the unit file" in block, "the active one is shown in full"


def test_the_block_is_bounded():
    plan = plans.build(
        title="x",
        phases=[{"title": f"Phase {n}", "tasks": [f"task {n} " + "y" * 200]} for n in range(8)],
    )
    assert len(plans.render_block(plan, budget=400)) < 500


def test_no_plan_is_an_empty_block():
    assert plans.render_block({}) == ""
    assert plans.render_block(None) == ""


# --- End to end -------------------------------------------------------------------
@pytest.fixture
def owner(db):
    """An administrator, because `tools.agent` is off by default and every test
    below is about what happens once agent chats are allowed at all."""
    user = User(name="Frodo", email="f@shire.test", password_hash=hash_password("x"))
    user.role = "admin"
    db.add(user)
    db.commit()
    return user


def _agent_chat(db, owner, mode=policy.MODE_EDIT, name="Box"):
    from lembas.db.models import SshProfile
    from lembas.services import settings_store

    profile = SshProfile(
        owner_id=owner.id, name=name, host="127.0.0.1", port=22, username="t",
        host_key="k", host_fingerprint="f", default_dir="/work",
    )
    connection = Connection(name=f"c-{name}", base_url="http://127.0.0.1:1", api_key_encrypted="")
    db.add_all([profile, connection])
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
    db.commit()
    settings_store.update(db, {"enabled": True}, key=settings_store.AGENTS)

    chat = Chat(
        user_id=owner.id, model_id="m", connection_id=connection.id, kind=KIND_AGENT,
        ssh_profile_id=profile.id, project_dir="/work", agent_mode=mode,
    )
    db.add(chat)
    db.commit()
    return chat


def _with_plan(db, chat, plan):
    message = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="", plan_json=plan)
    db.add(message)
    db.commit()
    chat.plan_message_id = message.id
    db.commit()
    return message


def test_plan_update_is_not_offered_without_a_plan(db, owner):
    """The skills asymmetry, avoided: a tool for changing something that does
    not exist costs a round to find out."""
    chat = _agent_chat(db, owner)
    names = set(tools_service.resolve_tools(db, chat, owner).by_name)

    assert "plan_update" not in names


def test_plan_update_is_offered_once_there_is_one(db, owner):
    chat = _agent_chat(db, owner)
    _with_plan(db, chat, V1)

    names = set(tools_service.resolve_tools(db, chat, owner).by_name)
    assert "plan_update" in names
    assert "plan_submit" not in names, "that one is Plan mode only"


def test_plan_submit_and_plan_update_are_never_offered_together(db, owner):
    chat = _agent_chat(db, owner, mode=policy.MODE_PLAN)
    _with_plan(db, chat, V1)

    names = set(tools_service.resolve_tools(db, chat, owner).by_name)
    assert "plan_submit" in names
    assert "plan_update" not in names


def test_the_plan_reaches_the_prompt(db, owner):
    """A plan the model cannot see is a plan it cannot keep current."""
    from lembas.services import harness

    chat = _agent_chat(db, owner)
    _with_plan(db, chat, V1)

    offered = tools_service.resolve_tools(db, chat, owner).schemas
    text = harness.compose(db, owner, offered, chat)

    assert "Tidy the logs" in text
    assert "Rotate it" in text
    assert "Keep it current" in text, "and the guidance to update it"


def test_no_plan_means_neither_the_section_nor_the_guidance(db, owner):
    from lembas.services import harness

    chat = _agent_chat(db, owner)
    offered = tools_service.resolve_tools(db, chat, owner).schemas
    text = harness.compose(db, owner, offered, chat)

    assert "The current plan" not in text
    assert "Keep it current" not in text


async def test_two_updates_in_one_reply_both_survive(db, owner):
    """The subtle one. A runner cannot write the message row -- `_persist` is
    the single writer -- so both updates would read the same stale plan from the
    database and the second would lose the first. They merge into the snapshot
    on AgentContext instead."""
    chat = _agent_chat(db, owner)
    _with_plan(db, chat, V1)
    resolved = tools_service.resolve_tools(db, chat, owner)
    context = tools_service.context_for(db, owner, chat, tools=resolved)

    await tools_service.run_tool(
        context, "plan_update", _json.dumps({"task_status": [{"id": "t1", "status": "done"}]})
    )
    second = await tools_service.run_tool(
        context, "plan_update", _json.dumps({"task_status": [{"id": "t2", "status": "done"}]})
    )

    tasks = {t["id"]: t["status"] for p in second.event["plan"]["phases"] for t in p["tasks"]}
    assert tasks["t1"] == "done", "the first update was not lost"
    assert tasks["t2"] == "done"


async def test_plan_update_never_ends_the_turn(db, owner):
    """`plan_submit` withdraws the tools for one final round because it ends the
    reply. Doing that here would stop the work dead every time a task was
    ticked off."""
    chat = _agent_chat(db, owner)
    _with_plan(db, chat, V1)
    resolved = tools_service.resolve_tools(db, chat, owner)
    context = tools_service.context_for(db, owner, chat, tools=resolved)

    outcome = await tools_service.run_tool(
        context, "plan_update", _json.dumps({"task_status": [{"id": "t1", "status": "done"}]})
    )
    assert not outcome.event.get("plan_final")


def test_plan_update_is_read_risk_so_it_does_not_ask(db, owner):
    """Otherwise carrying out a four-task plan means four approval cards, each
    approving a bookkeeping entry. Recorded as a decision, not an accident."""
    chat = _agent_chat(db, owner)
    _with_plan(db, chat, V1)
    resolved = tools_service.resolve_tools(db, chat, owner)

    assert resolved.by_name["plan_update"].risk == tools_service.RISK_READ
    decision = policy.decide(
        mode=policy.MODE_MANUAL, risk=tools_service.RISK_READ, tool_name="plan_update"
    )
    # Manual still asks about everything, which is what Manual means. Edit and
    # Auto -- where the work is actually carried out -- do not.
    assert decision.verdict == policy.ASK
    for mode in (policy.MODE_EDIT, policy.MODE_AUTO):
        assert (
            policy.decide(mode=mode, risk=tools_service.RISK_READ, tool_name="plan_update").verdict
            == policy.ALLOW
        )


def test_the_context_carries_the_plan(db, owner):
    chat = _agent_chat(db, owner)
    _with_plan(db, chat, V1)

    context = session.resolve(db, chat, owner)
    assert context.plan["title"] == "Tidy the logs"


def test_a_plan_pointer_at_another_chats_message_is_ignored(db, owner):
    """A plain id, not a foreign key, so it is validated on read."""
    chat = _agent_chat(db, owner)
    other = _agent_chat(db, owner, name="Other")
    message = _with_plan(db, other, V1)
    chat.plan_message_id = message.id
    db.commit()

    assert session.resolve(db, chat, owner).plan == {}


# --- Being asked to carry on ------------------------------------------------------
def _stub(rounds, seen):
    async def stream_chat(_endpoint, payload):
        seen.append(payload)
        for chunk in rounds[min(len(seen) - 1, len(rounds) - 1)]:
            yield chunk

    return stream_chat


def _text(text):
    return {"choices": [{"delta": {"content": text}}]}


def _call(name, arguments):
    return {"choices": [{"delta": {"tool_calls": [
        {"index": 0, "id": "c1", "function": {"name": name, "arguments": arguments}}]}}]}


def _reply(db, chat):
    from lembas.services import generation as generation_service

    db.add(Message(chat_id=chat.id, role="user", content="do it", complete=True))
    db.commit()
    assistant = Message(chat_id=chat.id, role=ROLE_ASSISTANT, content="", complete=False)
    db.add(assistant)
    db.commit()
    return generation_service.Generation(chat_id=chat.id, message_id=assistant.id)


async def _run(monkeypatch, generation, rounds, seen):
    from lembas.services import generation as generation_service

    monkeypatch.setattr(generation_service, "stream_chat", _stub(rounds, seen))

    async def _no_title(*_a, **_k):
        return ""

    monkeypatch.setattr("lembas.services.chat.generate_title", _no_title)
    await generation_service._run(generation)


async def test_stopping_with_open_tasks_is_answered_with_carry_on(db, owner, monkeypatch):
    """The half prompting cannot do. core.keep_working tells it not to stop
    halfway; this is what happens when it does anyway."""
    chat = _agent_chat(db, owner, mode=policy.MODE_AUTO)
    _with_plan(db, chat, V1)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(monkeypatch, generation, [[_text("I have done the first bit.")],
                                         [_text("All finished.")]], seen)

    # Three requests: the reply, then a nudge, then a second nudge -- this
    # model never marks anything done, so the plan stays open and it is asked
    # until MAX_NUDGES runs out. That it gives up is the next test.
    assert len(seen) == 3, "it was asked again"
    nudge = seen[1]["messages"][-1]
    assert nudge["role"] == "user"
    assert "still has work in it" in nudge["content"]
    assert "Rotate it" in nudge["content"], "and says which tasks"
    # Its own words go back with it, or it is asked to carry on from a
    # transcript in which it never spoke.
    assert seen[1]["messages"][-2]["content"] == "I have done the first bit."


async def test_a_finished_plan_is_believed(db, owner, monkeypatch):
    chat = _agent_chat(db, owner, mode=policy.MODE_AUTO)
    done = plans.build(title="x", phases=[{"title": "P", "tasks": ["a"]}])
    done["phases"][0]["tasks"][0]["status"] = "done"
    _with_plan(db, chat, done)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(monkeypatch, generation, [[_text("All done.")]], seen)

    assert len(seen) == 1


async def test_a_chat_with_no_plan_is_never_nudged(db, owner, monkeypatch):
    """There is nothing to be objectively wrong about, so a model that says it
    has finished is believed."""
    chat = _agent_chat(db, owner, mode=policy.MODE_AUTO)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(monkeypatch, generation, [[_text("All done.")]], seen)

    assert len(seen) == 1


async def test_plan_mode_is_never_nudged(db, owner, monkeypatch):
    """plan_submit ends the turn deliberately. Nudging past it would argue with
    the whole point of the mode."""
    chat = _agent_chat(db, owner, mode=policy.MODE_PLAN)
    _with_plan(db, chat, V1)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(monkeypatch, generation, [[_text("Here is what I would do.")]], seen)

    assert len(seen) == 1


async def test_it_gives_up_after_two_and_says_so(db, owner, monkeypatch):
    """A model that has nothing left to do must be able to say so and be
    believed rather than argued with indefinitely -- and a reply that stopped
    twice with work outstanding is worth being able to see afterwards."""
    chat = _agent_chat(db, owner, mode=policy.MODE_AUTO)
    _with_plan(db, chat, V1)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(monkeypatch, generation, [[_text("Nothing more from me.")]], seen)

    from lembas.services import generation as generation_service

    assert len(seen) == generation_service.MAX_NUDGES + 1
    assert generation.tool_events[-1]["status"] == "error"
    assert "asked twice" in generation.tool_events[-1]["error"]


async def test_calling_a_tool_again_resets_the_count(db, owner, monkeypatch):
    """The count is of consecutive stops. A model that stops, is nudged, does
    some work and stops again has not run out of patience."""
    chat = _agent_chat(db, owner, mode=policy.MODE_AUTO)
    _with_plan(db, chat, V1)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(
        monkeypatch,
        generation,
        [[_text("Pausing.")], [_call("file_list", '{"path": "."}')], [_text("Pausing again.")]],
        seen,
    )

    # Without the reset this run would end after two nudges, at three requests.
    # The tool call in the middle clears the count, so it gets more than that --
    # which is the property, and does not depend on where the stub stops.
    assert len(seen) > 3


async def test_the_switch_turns_it_off(db, owner, monkeypatch):
    from lembas.services import settings_store

    settings_store.update(db, {"nudge_unfinished": False}, key=settings_store.AGENTS)
    chat = _agent_chat(db, owner, mode=policy.MODE_AUTO)
    _with_plan(db, chat, V1)
    generation = _reply(db, chat)

    seen: list[dict] = []
    await _run(monkeypatch, generation, [[_text("All done.")]], seen)

    assert len(seen) == 1