LLeMbas/tests/test_agent_draft.py

"""The terminal and the canvas before a chat exists.

Chats are created lazily and that rule is kept: a draft is not a chat and never
becomes one. What it does is hold an id, three facts and a tab strip, so the
panels can work on the screen where you are still deciding which machine to work
on -- and hand all of it over when the first prompt creates the real chat.

The security-shaped tests here are the two refusals. `canvas._load_file`
authorises with `attachment.chat_id != chat.id`, and a draft upload is stored
with `chat_id=None`; a draft whose chat carried no id would make that comparison
`None != None`, which is False.
"""

from __future__ import annotations

import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select

from lembas.db.models import KIND_AGENT, Chat, Connection, Model, SshProfile, User
from lembas.services import settings_store
from lembas.services.agent import draft as draft_service
from lembas.services.agent import terminal as terminal_service


@pytest.fixture(autouse=True)
def _clean():
    draft_service.clear()
    yield
    draft_service.clear()


@pytest.fixture
def target(db, registered):
    settings_store.update(
        db, {"enabled": True, "terminal_enabled": True}, key=settings_store.AGENTS
    )
    user = db.scalar(select(User))
    profile = SshProfile(
        owner_id=user.id,
        name="Box",
        host="127.0.0.1",
        port=1,
        username="nobody",
        host_key="ssh-ed25519 AAAA",
        host_fingerprint="SHA256:x",
        default_dir="/srv/project",
    )
    connection = Connection(name="c", 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()
    return user, profile


# --- The id --------------------------------------------------------------------
def test_the_same_target_is_the_same_draft():
    """Derived rather than invented, so coming back to the new-chat screen finds
    the shell already running there instead of quietly opening a second."""
    one = draft_service.remember("u1", "p1", "/srv/x")
    two = draft_service.remember("u1", "p1", "/srv/x")

    assert one.id == two.id
    assert draft_service.is_draft(one.id)


def test_a_different_directory_or_owner_is_a_different_draft():
    base = draft_service.remember("u1", "p1", "/srv/x").id

    assert draft_service.remember("u1", "p1", "/srv/y").id != base
    assert draft_service.remember("u1", "p2", "/srv/x").id != base
    # Two people pointed at the same directory of the same connection would
    # otherwise share a *shell*, and one chat one shell is scoped to a person.
    assert draft_service.remember("u2", "p1", "/srv/x").id != base


def test_a_draft_id_cannot_be_confused_with_a_chat_id():
    """`Chat.id` is 32 hex characters. Nothing here can collide with one."""
    assert not draft_service.is_draft("a" * 32)
    assert not draft_service.is_draft("")


def test_somebody_elses_draft_is_not_readable():
    """The id is a hash of the owner, so it cannot be guessed -- but that is not
    an authorisation, and the next caller might build the id differently."""
    made = draft_service.remember("u1", "p1", "/srv/x")

    assert draft_service.get(made.id, "u1") is not None
    assert draft_service.get(made.id, "u2") is None


# --- The transient chat --------------------------------------------------------
def test_the_transient_chat_carries_what_the_panels_read():
    """`agent_ready`, `_executor`, `_load_agent` and `agent_session.resolve` read
    exactly these four, and none of them queries or writes the row -- which is
    why none of them had to learn what a draft is."""
    made = draft_service.remember("u1", "p1", "/srv/x")
    chat = draft_service.as_chat(made)

    assert chat.user_id == "u1"
    assert chat.kind == KIND_AGENT
    assert chat.ssh_profile_id == "p1"
    assert chat.project_dir == "/srv/x"
    assert chat.canvas_json == {}


def test_the_transient_chat_has_an_id_and_is_never_saved(db):
    """The id matters: `_load_file` compares it against an attachment's, and a
    draft upload has `chat_id=None`. `None != None` is False."""
    chat = draft_service.as_chat(draft_service.remember("u1", "p1", "/srv/x"))

    assert chat.id, "a column default is applied at flush, and this is never flushed"
    assert db.get(Chat, chat.id) is None, "and it must not have reached the database"


# --- The refusals --------------------------------------------------------------
def test_the_two_sources_that_need_a_row_are_refused_by_name():
    """Stated rather than left to fall out of an id comparison. `scratch` would
    write a row keyed on a chat that does not exist; `file` is the hole."""
    assert draft_service.refuses("scratch")
    assert draft_service.refuses("file")
    assert not draft_service.refuses("agent")
    assert not draft_service.refuses("note")


def test_opening_a_refused_source_says_so_rather_than_failing(client, db, target):
    user, profile = target
    made = draft_service.remember(user.id, profile.id, "/srv/project")

    response = client.post(
        f"/api/chats/{made.id}/canvas/tabs", data={"key": f"scratch:{made.id}"}
    )

    assert response.status_code == 200
    assert "once this chat exists" in response.text
    assert db.query(Chat).count() == 0, "and no row was created on the way"


# --- Adoption ------------------------------------------------------------------
def test_sending_the_first_prompt_adopts_the_shell(client, db, target, monkeypatch):
    """The shell opened while deciding becomes the chat's shell, scrollback and
    all. A re-key, not a reconnect: the browser navigates after `start_chat` and
    attaches to the session now living under the real id."""
    user, profile = target
    made = draft_service.remember(user.id, profile.id, "/srv/project")

    session = _fake_session(made.id, user.id, profile.id, "/srv/project")
    terminal_service._SESSIONS[made.id] = session
    try:
        client.post(
            "/api/chats/start",
            data={
                "content": "Do the thing.",
                "kind": "agent",
                "ssh_profile_id": profile.id,
                "project_dir": "/srv/project",
                "model_id": "m",
                "draft_id": made.id,
            },
        )
        chat = db.scalar(select(Chat))

        assert terminal_service._SESSIONS.get(chat.id) is session
        assert made.id not in terminal_service._SESSIONS
        # Both, or `close_for_profile` and the reaper -- which pop by the field
        # rather than by the key -- would leave a dead session findable.
        assert session.chat_id == chat.id
    finally:
        terminal_service._SESSIONS.clear()


def test_a_shell_on_a_different_target_is_not_transplanted(client, db, target):
    """`_new_chat` settles the directory last: an empty one falls back to the
    connection's own. A shell opened elsewhere belongs to the draft it was
    opened under, and is reaped on idle rather than moved."""
    user, profile = target
    made = draft_service.remember(user.id, profile.id, "/somewhere/else")

    session = _fake_session(made.id, user.id, profile.id, "/somewhere/else")
    terminal_service._SESSIONS[made.id] = session
    try:
        client.post(
            "/api/chats/start",
            data={
                "content": "Do the thing.",
                "kind": "agent",
                "ssh_profile_id": profile.id,
                "project_dir": "/srv/project",
                "model_id": "m",
                "draft_id": made.id,
            },
        )
        chat = db.scalar(select(Chat))

        assert chat.id not in terminal_service._SESSIONS
        assert session.chat_id == made.id
    finally:
        terminal_service._SESSIONS.clear()


def test_open_tabs_are_carried_onto_the_chat(client, db, target):
    user, profile = target
    made = draft_service.remember(user.id, profile.id, "/srv/project")
    made.canvas_json = {
        "tabs": [
            {"key": "agent:/srv/project/main.py", "title": "main.py", "source": "agent"},
            # Dropped rather than carried across to fail on first click.
            {"key": "scratch:whatever", "title": "Scratch", "source": "scratch"},
        ],
        "active": "agent:/srv/project/main.py",
    }

    client.post(
        "/api/chats/start",
        data={
            "content": "Do the thing.",
            "kind": "agent",
            "ssh_profile_id": profile.id,
            "project_dir": "/srv/project",
            "model_id": "m",
            "draft_id": made.id,
        },
    )
    chat = db.scalar(select(Chat))

    assert [t["key"] for t in chat.canvas_json["tabs"]] == ["agent:/srv/project/main.py"]
    assert draft_service.get(made.id, user.id) is None, "and the draft is done with"


def test_a_draft_belonging_to_somebody_else_adopts_nothing(client, db, target):
    user, profile = target
    theirs = draft_service.remember("someone-else", profile.id, "/srv/project")
    theirs.canvas_json = {"tabs": [{"key": "agent:/etc/shadow", "source": "agent"}]}

    client.post(
        "/api/chats/start",
        data={
            "content": "Do the thing.",
            "kind": "agent",
            "ssh_profile_id": profile.id,
            "project_dir": "/srv/project",
            "model_id": "m",
            "draft_id": theirs.id,
        },
    )
    chat = db.scalar(select(Chat))

    assert not (chat.canvas_json or {}).get("tabs")


# --- The screen ----------------------------------------------------------------
def test_the_new_chat_screen_offers_both_panels(client: TestClient, target):
    """The decision this reverses: they used to require a chat, which meant they
    were absent on the one screen where you are choosing a machine."""
    html = client.get("/chat").text

    assert "data-terminal" in html
    assert "data-canvas" in html
    assert "js/draft.js" in html


def test_the_draft_route_answers_with_a_stable_id(client: TestClient, target):
    _user, profile = target

    first = client.get(f"/api/agents/{profile.id}/draft?dir=/srv/project").json()
    again = client.get(f"/api/agents/{profile.id}/draft?dir=/srv/project").json()

    assert first["id"] == again["id"]
    assert draft_service.is_draft(first["id"])
    assert first["dir"] == "/srv/project"


def test_the_draft_route_refuses_somebody_elses_connection(client: TestClient, db, target):
    _user, profile = target
    stranger = User(email="s@x.test", name="S", password_hash="x")
    db.add(stranger)
    db.commit()
    profile.owner_id = stranger.id
    db.commit()

    assert client.get(f"/api/agents/{profile.id}/draft").status_code == 404


class _FakeSession:
    def __init__(self, chat_id, owner_id, profile_id, project_dir):
        self.chat_id = chat_id
        self.owner_id = owner_id
        self.profile_id = profile_id
        self.project_dir = project_dir
        self.closed = False


def _fake_session(chat_id, owner_id, profile_id, project_dir):
    return _FakeSession(chat_id, owner_id, profile_id, project_dir)


# --- Against a real machine ----------------------------------------------------
asyncssh = pytest.importorskip("asyncssh")


async def test_a_draft_canvas_opens_and_saves_a_project_file(db, registered, tmp_path):
    """The point of the whole thing: look around the machine, and edit something
    on it, before committing to a conversation about it.

    Against a real sshd and through `canvas_service` directly rather than the
    TestClient -- the client drives the app on another thread, and the server
    here is bound to this test's loop. What is being claimed is that a
    *transient* chat drives the same SFTP path a real one does, and that is what
    this exercises.
    """
    from lembas.services import canvas as canvas_service
    from lembas.services.agent import ssh as ssh_service
    from tests.test_canvas_ssh import _Server

    project = tmp_path / "project"
    project.mkdir()
    (project / "main.py").write_text("print('before')\n")

    server = await asyncssh.create_server(
        _Server,
        "127.0.0.1",
        0,
        server_host_keys=[asyncssh.generate_private_key("ssh-ed25519")],
        # SFTP only: the canvas never goes through a shell -- a path is a path.
        sftp_factory=True,
    )
    port = next(iter(server.sockets)).getsockname()[1]
    line, fingerprint = await ssh_service.capture_host_key("127.0.0.1", port)
    try:
        settings_store.update(db, {"enabled": True}, key=settings_store.AGENTS)
        user = db.scalar(select(User))
        profile = SshProfile(
            owner_id=user.id,
            name="Box",
            host="127.0.0.1",
            port=port,
            username="tester",
            host_key=line,
            host_fingerprint=fingerprint,
            default_dir=str(project),
        )
        db.add(profile)
        db.commit()

        made = draft_service.remember(user.id, profile.id, str(project))
        chat = draft_service.as_chat(made)
        key = f"agent:{project}/main.py"

        doc = await canvas_service.load(db, user, chat, key)
        assert "before" in doc.text

        await canvas_service.save(db, user, chat, key, "print('after')\n", doc.revision)
        assert (project / "main.py").read_text() == "print('after')\n"

        # And none of it reached the database.
        assert db.query(Chat).count() == 0
    finally:
        server.close()
        await server.wait_closed()