LLeMbas/tests/test_agent_transcript.py

"""An agent reply, rendered end to end.

The steps rewrite touched the one template every bubble goes through and the
frames that fill it, and an agent reply is the case with all of it at once:
reasoning, prose, tool events, a plan card, a project listing. `test_steps.py`
covers the builder as pure functions; this covers the page.
"""

from __future__ import annotations

import json as _json

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

from lembas.db.models import KIND_AGENT, Chat, Connection, Message, Model, SshProfile, User
from lembas.services import generation as generation_service
from lembas.services import settings_store


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


def _tool_chunk(name: str, arguments: str) -> dict:
    return {
        "choices": [
            {
                "delta": {
                    "tool_calls": [
                        {
                            "index": 0,
                            "id": "call_1",
                            "function": {"name": name, "arguments": arguments},
                        }
                    ]
                }
            }
        ]
    }


def _stub_stream(rounds):
    calls = {"n": 0}

    async def stream(_endpoint, _payload):
        index = min(calls["n"], len(rounds) - 1)
        calls["n"] += 1
        for chunk in rounds[index]:
            yield chunk

    return stream


@pytest.fixture
def agent_chat(db, registered):
    """An agent chat whose connection exists but has no reachable host.

    Deliberately unreachable: what is under test is the *rendering* of a reply
    that called a tool, and a tool that fails still produces an event, which is
    the thing the transcript has to lay out.
    """
    settings_store.update(
        db,
        {"enabled": True, "background_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()

    chat = Chat(
        user_id=user.id,
        model_id="m",
        connection_id=connection.id,
        kind=KIND_AGENT,
        ssh_profile_id=profile.id,
        project_dir="/srv/project",
        agent_mode="auto",
    )
    db.add(chat)
    db.commit()
    return chat


async def test_an_agent_reply_that_called_a_tool_renders(
    client: TestClient, db, agent_chat, registered, monkeypatch
):
    """The regression this file exists for: the bubble is rendered from the
    steps now, and an agent reply is where every kind of step appears at once."""
    db.add(Message(chat_id=agent_chat.id, role="user", content="Do the thing.", complete=True))
    reply = Message(chat_id=agent_chat.id, role="assistant", content="", complete=False)
    db.add(reply)
    db.commit()
    reply_id = reply.id

    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [
                    _text_chunk("Listing the project first. "),
                    _tool_chunk("file_list", _json.dumps({"path": "."})),
                ],
                [_text_chunk("Nothing I can reach.")],
            ]
        ),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", lambda *a, **k: "")

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

    db.expire_all()
    stored = db.get(Message, reply_id)
    assert stored.tool_calls_json, "the tool event was recorded"
    assert stored.steps_json, "and so were the marks"

    page = client.get(f"/chat/{agent_chat.id}")
    assert page.status_code == 200
    assert "Listing the project first" in page.text
    assert "Nothing I can reach" in page.text
    assert page.text.index("Listing the project first") < page.text.index("tool-activity")


def test_the_page_renders_before_any_reply_exists(client: TestClient, agent_chat, registered):
    """An agent chat with nothing in it. The composer takes a different branch
    here -- the mode select, the jobs chip -- and it is the screen somebody sees
    first."""
    page = client.get(f"/chat/{agent_chat.id}")

    assert page.status_code == 200
    assert 'name="agent_mode"' in page.text


def test_starting_an_agent_chat_from_the_composer_renders(
    client: TestClient, db, registered, agent_chat
):
    """`POST /api/chats/start` returns the thread it just created. This is the
    exact moment the reader reported a blank screen."""
    profile = db.scalar(select(SshProfile))
    response = client.post(
        "/api/chats/start",
        data={
            "content": "Do the thing.",
            "kind": "agent",
            "ssh_profile_id": profile.id,
            "project_dir": "/srv/project",
            "model_id": "m",
        },
    )

    # 204 plus HX-Redirect is the contract: the row is created here and the
    # browser then navigates to it. The blank screen is on the page it lands on.
    assert response.status_code == 204
    assert response.headers["HX-Redirect"].startswith("/chat/")

    landed = client.get(response.headers["HX-Redirect"])
    assert landed.status_code == 200
    assert "Do the thing." in landed.text


def test_the_page_renders_while_the_reply_is_still_unfinished(
    client: TestClient, db, agent_chat, registered
):
    """The streaming branch of the bubble, which is what somebody sees for the
    whole of a long agent reply -- and what they saw blank.

    An unfinished assistant message is the *only* thing that starts a
    generation, so this is also the page that sets the reply going.
    """
    db.add(Message(chat_id=agent_chat.id, role="user", content="Do the thing.", complete=True))
    db.add(Message(chat_id=agent_chat.id, role="assistant", content="", complete=False))
    db.commit()

    page = client.get(f"/chat/{agent_chat.id}")

    assert page.status_code == 200
    assert "Do the thing." in page.text
    assert 'sse-swap="steps"' in page.text, "the container the frames land in"
    assert 'sse-swap="render"' in page.text


async def test_the_stream_delivers_the_steps_frame(
    client: TestClient, db, agent_chat, registered, monkeypatch
):
    """The SSE path with the new frame set, which is the one place the steps are
    rendered without a `Message` row to hang them on."""
    db.add(Message(chat_id=agent_chat.id, role="user", content="Do it.", complete=True))
    reply = Message(chat_id=agent_chat.id, role="assistant", content="", complete=False)
    db.add(reply)
    db.commit()
    reply_id = reply.id

    monkeypatch.setattr(
        generation_service,
        "stream_chat",
        _stub_stream(
            [
                [
                    _text_chunk("Listing first. "),
                    _tool_chunk("file_list", _json.dumps({"path": "."})),
                ],
                [_text_chunk("Done.")],
            ]
        ),
    )
    monkeypatch.setattr("lembas.services.chat.generate_title", lambda *a, **k: "")

    with client.stream(
        "GET", f"/api/chats/{agent_chat.id}/messages/{reply_id}/stream"
    ) as response:
        assert response.status_code == 200
        body = "".join(response.iter_lines())

    assert "event: steps" in body, "the closed steps never reached the browser"
    assert "event: done" in body
    assert "Listing first" in body