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 == {}