LLeMbas/tests/test_agent_tools.py

"""The agent tools: when they are offered, and what stops them running.

Everything here goes through a real SSH server on 127.0.0.1, so the gate and the
approval loop are exercised against something that genuinely executes rather
than a stub that always agrees.
"""

from __future__ import annotations

import asyncio

import pytest

from lembas.db.models import (
    KIND_AGENT,
    ROLE_ASSISTANT,
    Chat,
    Connection,
    Message,
    Model,
    SshProfile,
    User,
)
from lembas.services import generation as generation_service
from lembas.services import interaction, settings_store
from lembas.services import tools as tools_service
from lembas.services.agent import policy, session
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):
    command = process.command or ""
    process.stdout.write(f"ran: {command}\n")
    process.exit(0)


@pytest.fixture
async def machine(tmp_path):
    """A real sshd, and a profile pointing at it with its key already pinned."""
    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 _setup(db, user_id, machine, *, mode=policy.MODE_MANUAL, enabled=True, kind=KIND_AGENT):
    """An agent chat pointed at the machine, with the feature switched on."""
    settings_store.update(db, {"enabled": enabled}, 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,
        ssh_profile_id=profile.id,
        project_dir=machine["dir"],
        agent_mode=mode,
    )
    db.add(chat)
    db.commit()
    return chat, profile


def _offered(db, chat, user) -> set[str]:
    return set(tools_service.resolve_tools(db, chat, user).by_name)


# --- The gate --------------------------------------------------------------------
async def test_the_agent_tools_are_offered_to_an_agent_chat(db, user_id, machine):
    chat, _profile = _setup(db, user_id, machine)
    names = _offered(db, chat, db.get(User, user_id))

    assert {"shell_run", "file_read", "file_write", "file_list"} <= names


async def test_an_ordinary_chat_gets_none_of_them(db, user_id, machine):
    """A plain conversation can never shell out, whatever it is asked."""
    chat, _profile = _setup(db, user_id, machine, kind="chat")
    assert "shell_run" not in _offered(db, chat, db.get(User, user_id))


async def test_nothing_is_offered_while_the_feature_is_off(db, user_id, machine):
    chat, _profile = _setup(db, user_id, machine, enabled=False)
    assert "shell_run" not in _offered(db, chat, db.get(User, user_id))


async def test_nothing_is_offered_without_the_permission(db, user_id, machine):
    chat, _profile = _setup(db, user_id, machine)
    user = db.get(User, user_id)
    user.role = "user"  # administrators pass everything
    settings_store.update(db, {"default_permissions": {"tools.agent": False}})
    db.commit()

    assert "shell_run" not in _offered(db, chat, user)


async def test_nothing_is_offered_without_the_model_capability(db, user_id, machine):
    chat, _profile = _setup(db, user_id, machine)
    model = db.scalar(tools_service.select(Model))
    model.capabilities_json = {"tools": True, "tool_agent": False}
    db.commit()

    assert "shell_run" not in _offered(db, chat, db.get(User, user_id))


async def test_a_disabled_connection_takes_the_tools_away(db, user_id, machine):
    """Offering a tool that fails on its first call is worse than not offering
    it, so every "no" collapses to an empty list."""
    chat, profile = _setup(db, user_id, machine)
    profile.enabled = False
    db.commit()

    assert "shell_run" not in _offered(db, chat, db.get(User, user_id))


async def test_somebody_elses_connection_is_not_reachable(db, user_id, machine):
    from lembas.security.passwords import hash_password

    chat, profile = _setup(db, user_id, machine)
    intruder = User(name="Sam", email="s@shire.test", password_hash=hash_password("x"))
    intruder.role = "user"
    db.add(intruder)
    db.commit()

    assert session.resolve(db, chat, intruder) is None


# --- Running something -----------------------------------------------------------
async def test_a_command_runs_on_the_machine(db, user_id, machine):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
    user = db.get(User, user_id)
    resolved = tools_service.resolve_tools(db, chat, user)
    context = tools_service.context_for(db, user, chat, tools=resolved)

    outcome = await tools_service.run_tool(
        context, "shell_run", '{"command": "echo hello"}'
    )
    assert outcome.event["status"] == "ok"
    assert "echo hello" in outcome.content


async def test_files_are_written_and_read_back(db, user_id, machine, tmp_path):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
    user = db.get(User, user_id)
    resolved = tools_service.resolve_tools(db, chat, user)
    context = tools_service.context_for(db, user, chat, tools=resolved)

    await tools_service.run_tool(
        context, "file_write", '{"path": "note.txt", "content": "a mallorn tree"}'
    )
    assert (tmp_path / "project" / "note.txt").read_text() == "a mallorn tree"

    read = await tools_service.run_tool(context, "file_read", '{"path": "note.txt"}')
    assert "mallorn" in read.content

    listed = await tools_service.run_tool(context, "file_list", "{}")
    assert "note.txt" in listed.content


# --- The runner backstop ----------------------------------------------------------
async def test_a_runner_refuses_what_the_mode_forbids(db, user_id, machine):
    """`_authorise` is the real gate and runs first. This is the belt to that
    brace: a call arriving by some other path -- a retry, a re-run button --
    must not walk past it."""
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
    user = db.get(User, user_id)
    resolved = tools_service.resolve_tools(db, chat, user)
    context = tools_service.context_for(db, user, chat, tools=resolved)

    outcome = await tools_service.run_tool(context, "shell_run", '{"command": "rm -rf /"}')
    assert outcome.event["status"] == "error"
    assert "not allowed" in outcome.content.lower()


async def test_plan_mode_still_reads(db, user_id, machine, tmp_path):
    (tmp_path / "project" / "readme.txt").write_text("contents")
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_PLAN)
    user = db.get(User, user_id)
    resolved = tools_service.resolve_tools(db, chat, user)
    context = tools_service.context_for(db, user, chat, tools=resolved)

    outcome = await tools_service.run_tool(context, "file_read", '{"path": "readme.txt"}')
    assert outcome.event["status"] == "ok"
    assert "contents" in outcome.content


async def test_edit_mode_writes_but_will_not_run(db, user_id, machine):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_EDIT)
    user = db.get(User, user_id)
    resolved = tools_service.resolve_tools(db, chat, user)
    context = tools_service.context_for(db, user, chat, tools=resolved)

    wrote = await tools_service.run_tool(
        context, "file_write", '{"path": "x.txt", "content": "hi"}'
    )
    assert wrote.event["status"] == "ok"

    ran = await tools_service.run_tool(context, "shell_run", '{"command": "ls"}')
    assert ran.event["status"] == "error"


# --- The approval card in the loop -------------------------------------------------
def _chunk(name: str, arguments: str, *, call_id="c1", index=0):
    return {
        "choices": [
            {
                "delta": {
                    "tool_calls": [
                        {
                            "index": index,
                            "id": call_id,
                            "function": {"name": name, "arguments": arguments},
                        }
                    ]
                }
            }
        ]
    }


def _text(text: str) -> dict:
    return {"choices": [{"delta": {"content": text}}]}


def _stub_stream(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


async def _until_paused(generation, *, timeout: float = 3.0):
    deadline = asyncio.get_running_loop().time() + timeout
    while asyncio.get_running_loop().time() < deadline:
        if generation.pending is not None:
            return generation.pending
        await asyncio.sleep(0.01)
    raise AssertionError("the reply never paused for approval")


def _pending_reply(db, chat):
    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 assistant.id


async def test_a_command_waits_for_approval_and_the_card_shows_it(
    db, user_id, machine, monkeypatch
):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
    message_id = _pending_reply(db, chat)
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [[_chunk("shell_run", '{"command": "rm -rf /tmp/x"}')], [_text("Done.")]], []
        ),
    )

    generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
    task = asyncio.create_task(generation_service._run(generation))
    pending = await _until_paused(generation)

    assert pending.kind == interaction.KIND_APPROVAL
    item = pending.items[0]
    # The exact command, verbatim. A card that paraphrased it would be
    # approving something other than what runs.
    assert item.detail == "rm -rf /tmp/x"
    assert "Test box" in item.title
    assert "Manual" in item.reason

    pending.resolve(interaction.DENY)
    await task


async def test_denying_reaches_the_model_as_words_and_runs_nothing(
    db, user_id, machine, monkeypatch, tmp_path
):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
    message_id = _pending_reply(db, chat)
    payloads: list[dict] = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [_chunk("file_write", '{"path": "never.txt", "content": "nope"}')],
                [_text("Understood.")],
            ],
            payloads,
        ),
    )

    generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
    task = asyncio.create_task(generation_service._run(generation))
    pending = await _until_paused(generation)
    pending.resolve(interaction.DENY)
    await task

    turns = [m for m in payloads[1]["messages"] if m.get("role") == "tool"]
    assert "declined" in turns[0]["content"].lower()
    # And it says not to look for another way round, because a model reads
    # "not allowed" as "not allowed like that" otherwise.
    assert "another way" in turns[0]["content"]
    assert not (tmp_path / "project" / "never.txt").exists()


async def test_allowing_runs_it(db, user_id, machine, monkeypatch, tmp_path):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
    message_id = _pending_reply(db, chat)
    payloads: list[dict] = []
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [_chunk("file_write", '{"path": "yes.txt", "content": "written"}')],
                [_text("Done.")],
            ],
            payloads,
        ),
    )

    generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
    task = asyncio.create_task(generation_service._run(generation))
    pending = await _until_paused(generation)
    pending.resolve(interaction.ALLOW)
    await task

    assert (tmp_path / "project" / "yes.txt").read_text() == "written"


async def test_auto_mode_never_pauses(db, user_id, machine, monkeypatch, tmp_path):
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
    message_id = _pending_reply(db, chat)
    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [_chunk("file_write", '{"path": "auto.txt", "content": "no asking"}')],
                [_text("Done.")],
            ],
            [],
        ),
    )

    generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
    await asyncio.wait_for(generation_service._run(generation), timeout=10)

    assert generation.pending is None
    assert (tmp_path / "project" / "auto.txt").read_text() == "no asking"


async def test_the_credential_is_cleared_when_the_reply_ends(db, user_id, machine, monkeypatch):
    """A finished Generation lingers five minutes so late followers get the
    final frames. A private key should not linger with it."""
    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_AUTO)
    message_id = _pending_reply(db, chat)

    captured = {}
    original = tools_service.context_for

    def capture(*args, **kwargs):
        context = original(*args, **kwargs)
        captured["context"] = context
        return context

    monkeypatch.setattr(tools_service, "context_for", capture)
    monkeypatch.setattr(
        generation_service, "stream_chat", _stub_stream([[_text("Nothing to do.")]], [])
    )

    generation = generation_service.Generation(chat_id=chat.id, message_id=message_id)
    await generation_service._run(generation)

    assert captured["context"].agent is not None
    assert captured["context"].agent.spec == {}, "the decrypted credential is dropped"


# --- The harness has to be able to name the machine ---------------------------
def test_the_registry_maps_the_agent_tools_to_their_family(db):
    """`harness._families` maps an offered tool *name* back to a family to
    decide which fragments apply, and it has no chat to resolve against. Without
    the agent tools listed here, `shell_run` resolves to no family and an agent
    chat is told nothing about the machine it is working on -- the same omission
    that cost custom tools their guidance once already."""
    book = tools_service.registry(db)
    for name in ("shell_run", "file_read", "file_write", "file_list"):
        assert name in book, name
        assert book[name].family == "agent"


async def test_the_harness_says_where_and_under_what_rules(db, user_id, machine):
    from lembas.services import harness

    chat, _profile = _setup(db, user_id, machine, mode=policy.MODE_MANUAL)
    user = db.get(User, user_id)
    offered = tools_service.resolve_tools(db, chat, user).schemas
    text = harness.compose(db, user, offered, chat)

    assert "Test box" in text, "which machine"
    assert machine["dir"] in text, "which directory"
    assert "Manual" in text, "what the mode permits"
    # The single most likely cause of "the agent seems stupid": `cd build`
    # followed by `make` fails silently otherwise.
    assert "fresh shell" in text


async def test_an_ordinary_chat_is_told_none_of_it(db, user_id, machine):
    from lembas.services import harness

    chat, _profile = _setup(db, user_id, machine, kind="chat")
    user = db.get(User, user_id)
    offered = tools_service.resolve_tools(db, chat, user).schemas
    text = harness.compose(db, user, offered, chat)

    assert "Test box" not in text
    assert "fresh shell" not in text


async def test_an_agent_chat_is_told_its_real_round_budget(db, user_id, machine):
    """MAX_ROUNDS is three. An agent chat gets forty, and telling it three would
    be a false fact about its own budget on every turn."""
    from lembas.services import harness

    settings_store.update(db, {"max_steps": 25}, key=settings_store.AGENTS)
    chat, _profile = _setup(db, user_id, machine)
    user = db.get(User, user_id)
    values = harness.context_variables(
        db, user, tools_service.resolve_tools(db, chat, user).schemas, chat
    )
    assert values["max_rounds"] == "25"