LLeMbas/tests/test_chat.py

"""Chat, folders, and the streaming reply path against a mocked endpoint."""

from __future__ import annotations

import json

import httpx
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import func, select

from lembas.db.models import Chat, Connection, Folder, Message, Model
from lembas.services import chat as chat_service
from lembas.services.crypto import encrypt
from lembas.services.llm.openai_client import Endpoint, LLMError, delta_text, list_models
from lembas.services.sse import event


# --- SSE framing -------------------------------------------------------------
def test_sse_event_framing():
    assert event("token", "hello") == "event: token\ndata: hello\n\n"


def test_sse_splits_newlines_across_data_lines():
    """A payload with a newline must become several data: lines. Sending a raw
    newline truncates the event, which is what breaks the first code block a
    model emits."""
    assert event("token", "a\nb") == "event: token\ndata: a\ndata: b\n\n"


def test_sse_round_trips_through_the_browser_rejoin_rule():
    payload = "line one\nline two\n\nline four"
    framed = event("token", payload)
    body = framed.split("\n", 1)[1]
    rejoined = "\n".join(
        line.removeprefix("data: ") for line in body.split("\n") if line.startswith("data:")
    )
    assert rejoined == payload


# --- Delta parsing -----------------------------------------------------------
def test_delta_text_reads_the_normal_shape():
    assert delta_text({"choices": [{"delta": {"content": "hi"}}]}) == "hi"


def test_delta_text_handles_typed_content_parts():
    chunk = {"choices": [{"delta": {"content": [{"type": "text", "text": "hi"}]}}]}
    assert delta_text(chunk) == "hi"


@pytest.mark.parametrize(
    "chunk", [{}, {"choices": []}, {"choices": [{}]}, {"choices": [{"delta": {}}]}]
)
def test_delta_text_tolerates_junk(chunk):
    """Providers vary; an unexpected chunk shape must not kill a reply."""
    assert delta_text(chunk) == ""


# --- Endpoint URL handling ---------------------------------------------------
@pytest.mark.parametrize(
    ("base", "expected"),
    [
        ("http://host:1234", "http://host:1234/v1/models"),
        ("http://host:1234/v1", "http://host:1234/v1/models"),
        ("http://host:1234/", "http://host:1234/v1/models"),
    ],
)
def test_base_url_with_or_without_v1(base, expected):
    """Users should not have to guess which form is expected."""
    assert Endpoint(base_url=base.rstrip("/"), api_key="", extra_headers={}).url(
        "models"
    ) == expected


def test_no_authorization_header_without_a_key():
    """Local runners often reject an empty bearer token outright."""
    assert "Authorization" not in Endpoint("http://h", "", {}).headers()
    assert Endpoint("http://h", "k", {}).headers()["Authorization"] == "Bearer k"


# --- Model discovery ---------------------------------------------------------
async def test_list_models_accepts_the_bare_list_shape():
    """The spec says {"data": [...]}, but some servers return a bare list."""

    def handler(_request):
        return httpx.Response(200, json=[{"id": "a"}, {"id": "b"}])

    original = httpx.AsyncClient

    class Patched(original):
        def __init__(self, **kwargs):
            super().__init__(transport=httpx.MockTransport(handler), **kwargs)

    httpx.AsyncClient = Patched
    try:
        assert [m["id"] for m in await list_models(Endpoint("http://h", "", {}))] == ["a", "b"]
    finally:
        httpx.AsyncClient = original


async def test_list_models_reports_a_rejected_key_readably():
    def handler(_request):
        return httpx.Response(401, json={"error": {"message": "Incorrect API key."}})

    original = httpx.AsyncClient

    class Patched(original):
        def __init__(self, **kwargs):
            super().__init__(transport=httpx.MockTransport(handler), **kwargs)

    httpx.AsyncClient = Patched
    try:
        with pytest.raises(LLMError) as caught:
            await list_models(Endpoint("http://h", "bad", {}))
        assert "rejected" in caught.value.message
        assert "Incorrect API key." in caught.value.message
    finally:
        httpx.AsyncClient = original


# --- Titles ------------------------------------------------------------------
def test_fallback_title_keeps_a_short_message_intact():
    assert chat_service.fallback_title("What is lembas?") == "What is lembas?"


def test_fallback_title_trims_on_a_word_boundary():
    title = chat_service.fallback_title("word " * 60)
    assert len(title) <= chat_service.MAX_TITLE_LENGTH + 1
    assert title.endswith("…")


def test_fallback_title_of_nothing():
    assert chat_service.fallback_title("   ") == "New chat"


async def test_the_title_prompt_carries_the_exchange(mock_http):
    """The wording is a fragment an administrator can edit, so what reaches the
    endpoint has to be the substituted text, not the template."""
    seen: list[str] = []

    def handler(request: httpx.Request) -> httpx.Response:
        seen.append(json.loads(request.content)["messages"][0]["content"])
        return httpx.Response(200, json={"choices": [{"message": {"content": "A short name"}}]})

    mock_http(handler)
    endpoint = Endpoint("http://x.test", "", {})
    title = await chat_service.generate_title(
        endpoint,
        "m",
        "What is lembas?",
        "Waybread.",
        template="Name this: {{question}} / {{answer}} / {{nonsense}}",
    )

    assert title == "A short name"
    assert seen == ["Name this: What is lembas? / Waybread. / {{nonsense}}"]


async def test_an_empty_title_prompt_asks_no_model_at_all(mock_http):
    """Clearing the fragment is how auto-titling is turned off. It must not
    cost a request that is then thrown away."""

    def handler(request: httpx.Request) -> httpx.Response:  # pragma: no cover - must not run
        raise AssertionError("the endpoint was contacted")

    mock_http(handler)
    endpoint = Endpoint("http://x.test", "", {})
    title = await chat_service.generate_title(
        endpoint, "m", "What is lembas?", "Waybread.", template="   "
    )
    assert title == "What is lembas?"


# --- Chats and folders (through the API) -------------------------------------
def _add_connection(db) -> Connection:
    # Port 1 refuses connections, which is what the error-path test relies on.
    connection = Connection(
        name="Test", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt("")
    )
    db.add(connection)
    db.commit()
    db.add(Model(connection_id=connection.id, model_id="test-model"))
    db.commit()
    return connection


# --- Starting a chat ---------------------------------------------------------
def test_starting_a_chat_creates_it_and_redirects(client: TestClient, db, registered):
    _add_connection(db)
    response = client.post("/api/chats/start", data={"content": "Hello there"})
    assert response.status_code == 204
    assert response.headers["HX-Redirect"].startswith("/chat/")

    chat = db.scalar(select(Chat))
    assert chat.model_id == "test-model"
    messages = db.scalars(select(Message).order_by(Message.created_at)).all()
    assert [m.role for m in messages] == ["user", "assistant"]
    assert messages[0].content == "Hello there"


def test_starting_with_nothing_creates_no_chat(client: TestClient, db, registered):
    """The whole point of lazy creation: an abandoned composer leaves nothing."""
    _add_connection(db)
    assert client.post("/api/chats/start", data={"content": "  "}).status_code == 204
    assert db.scalar(select(Chat)) is None


def test_visiting_the_chat_page_creates_nothing(client: TestClient, db, registered):
    _add_connection(db)
    assert client.get("/chat").status_code == 200
    assert db.scalar(select(Chat)) is None


def test_starting_a_chat_honours_the_requested_model(client: TestClient, db, registered):
    """The pinned-model shortcuts pass ?model=, which arrives here."""
    connection = _add_connection(db)
    db.add(Model(connection_id=connection.id, model_id="other-model", position=5))
    db.commit()

    client.post("/api/chats/start", data={"content": "hi", "model_id": "other-model"})
    assert db.scalar(select(Chat)).model_id == "other-model"


def test_starting_a_chat_ignores_a_model_you_cannot_reach(client: TestClient, db, registered):
    connection = _add_connection(db)
    db.add(Model(connection_id=connection.id, model_id="secret", position=5, public=False))
    db.commit()

    client.post("/auth/logout", follow_redirects=False)
    client.post(
        "/auth/register",
        data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
        follow_redirects=False,
    )
    client.post("/api/chats/start", data={"content": "hi", "model_id": "secret"})
    assert db.scalar(select(Chat)).model_id == "test-model"


def test_posting_a_message_stores_both_turns(
    client: TestClient, db, registered, make_chat, monkeypatch
):
    """What the route itself does, with no reply running behind it.

    `ensure` is stubbed out because the generation is a background task: it
    would race this test to the database, and against a connection that
    refuses instantly it sometimes wins -- writing the error and marking the
    row complete before the assertions below can read it. What the route
    guarantees is the pair of rows and the streaming shell; whether a reply has
    got anywhere yet is a different test's business.
    """
    from lembas.services import generation as generation_service

    monkeypatch.setattr(generation_service, "ensure", lambda *a, **k: None)

    _add_connection(db)
    chat_id = make_chat()

    response = client.post(f"/api/chats/{chat_id}/messages", data={"content": "Hello there"})
    assert response.status_code == 200

    messages = db.scalars(select(Message).order_by(Message.created_at)).all()
    assert [m.role for m in messages] == ["user", "assistant"]
    assert messages[0].content == "Hello there"
    # The assistant row is created empty and incomplete; that is what carries
    # the sse-connect the browser uses to start the stream.
    assert messages[1].content == ""
    assert messages[1].complete is False
    assert "sse-connect" in response.text


def test_empty_message_is_ignored(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    assert client.post(f"/api/chats/{chat_id}/messages", data={"content": "   "}).status_code == 204
    assert db.scalar(select(Message)) is None


def test_a_chat_belonging_to_someone_else_is_not_found(
    client: TestClient, db, registered, make_chat
):
    _add_connection(db)
    chat_id = make_chat()

    client.post("/auth/logout", follow_redirects=False)
    client.post(
        "/auth/register",
        data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
        follow_redirects=False,
    )
    # 404 and not 403: whether that id exists is not this endpoint's to reveal.
    assert client.get(f"/chat/{chat_id}").status_code == 404


def test_renaming_a_chat_stops_it_being_auto_titled(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    client.patch(f"/api/chats/{chat_id}", data={"title": "My own title"})

    chat = db.get(Chat, chat_id)
    db.refresh(chat)
    assert chat.title == "My own title"
    assert chat.title_generated is True


def test_deleting_a_chat_removes_its_messages(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "Hello"})

    client.delete(f"/api/chats/{chat_id}")
    assert db.scalar(select(Chat)) is None
    assert db.scalar(select(Message)) is None


def test_deleting_a_folder_keeps_the_chats_inside_it(client: TestClient, db, registered, make_chat):
    """Losing a conversation to a mis-clicked folder delete is unforgivable."""
    _add_connection(db)
    client.post("/api/folders", data={"name": "Quests"})
    folder = db.scalar(select(Folder))

    chat_id = make_chat()
    client.patch(f"/api/chats/{chat_id}", data={"folder_id": folder.id})

    client.delete(f"/api/folders/{folder.id}")

    chat = db.get(Chat, chat_id)
    db.refresh(chat)
    assert chat is not None
    assert chat.folder_id is None


def test_an_archived_chat_inside_a_folder_is_not_listed(
    client: TestClient, db, registered, make_chat
):
    """Regression: the unfiled list has always filtered archived chats, but the
    folder branch went through the ORM relationship and filtered nothing, so an
    archived chat kept showing as long as it was filed."""
    _add_connection(db)
    client.post("/api/folders", data={"name": "Quests"})
    folder = db.scalar(select(Folder))

    chat_id = make_chat()
    client.patch(f"/api/chats/{chat_id}", data={"folder_id": folder.id, "title": "Mount Doom"})
    chat = db.get(Chat, chat_id)
    chat.archived = True
    db.commit()

    page = client.get("/chat").text
    assert "Mount Doom" not in page
    # And the folder must say so, rather than claiming to hold something.
    assert "Empty" in page


def test_a_folder_cannot_be_moved_inside_itself(client: TestClient, db, registered):
    client.post("/api/folders", data={"name": "Outer"})
    folder = db.scalar(select(Folder))
    response = client.patch(f"/api/folders/{folder.id}", data={"parent_id": folder.id})
    assert response.status_code == 400


# --- Request building --------------------------------------------------------
def test_request_forwards_only_known_sampling_parameters(db, user_id):
    """A stray key in params_json must not become a 400 from the provider that
    looks like a LLeMbas bug."""
    connection = _add_connection(db)
    chat = Chat(
        user_id=user_id,
        model_id="test-model",
        connection_id=connection.id,
        params_json={"temperature": 0.4, "nonsense": "drop me"},
    )
    db.add(chat)
    db.commit()

    payload = chat_service.build_request(db, chat)
    assert payload["temperature"] == 0.4
    assert "nonsense" not in payload


def test_history_skips_failed_and_empty_turns(db, user_id):
    """Sending an empty assistant message upsets several providers."""
    connection = _add_connection(db)
    chat = Chat(user_id=user_id, model_id="test-model", connection_id=connection.id)
    db.add(chat)
    db.commit()

    db.add_all(
        [
            Message(chat_id=chat.id, role="user", content="one"),
            Message(chat_id=chat.id, role="assistant", content="", error="boom"),
            Message(chat_id=chat.id, role="user", content="two"),
        ]
    )
    db.commit()

    messages = chat_service.build_request(db, chat)["messages"]
    contents = [m["content"] for m in messages if m["role"] != "system"]
    assert contents == ["one", "two"]


def test_system_prompt_leads_the_message_list(db, user_id):
    connection = _add_connection(db)
    chat = Chat(
        user_id=user_id,
        model_id="test-model",
        connection_id=connection.id,
        system_prompt="You are terse.",
    )
    db.add(chat)
    db.commit()

    messages = chat_service.build_request(db, chat)["messages"]
    assert messages[0]["role"] == "system"
    # The harness precedes it inside the same message; the authored prompt is
    # last, where it is closest to the conversation.
    assert messages[0]["content"].endswith("You are terse.")


def test_streaming_reports_an_unreachable_endpoint_in_the_thread(
    client: TestClient, db, registered
, make_chat):
    """A failed turn must never be an unexplained blank bubble."""
    _add_connection(db)  # points at 127.0.0.1:1, which refuses connections
    chat_id = make_chat()
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "Hello"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))

    response = client.get(f"/api/chats/{chat_id}/messages/{message.id}/stream")
    assert response.status_code == 200
    assert "Could not reach" in response.text
    assert "alert--error" in response.text

    db.refresh(message)
    assert message.complete is True
    assert message.error


# --- Stopping a stream -------------------------------------------------------
def test_stopping_asks_the_generation_to_stop(client: TestClient, db, registered, make_chat):
    """A half-written answer the reader chose to cut short is still worth
    having; discarding it would be a surprise."""
    from lembas.services import generation as generation_service

    _add_connection(db)
    chat_id = make_chat()
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))

    assert client.post(
        f"/api/chats/{chat_id}/messages/{message.id}/stop"
    ).status_code == 204

    running = generation_service.get(message.id)
    # The endpoint points at 127.0.0.1:1, so the task may already have failed
    # and finished; either way the request must be accepted, not error.
    assert running is None or running.cancel or running.done


def test_stopping_someone_elses_message_is_refused(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))

    client.post("/auth/logout", follow_redirects=False)
    client.post(
        "/auth/register",
        data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
        follow_redirects=False,
    )
    assert client.post(
        f"/api/chats/{chat_id}/messages/{message.id}/stop"
    ).status_code == 404


def test_the_streaming_bubble_carries_the_sse_connection(
    client: TestClient, db, registered, make_chat
):
    """Stop lives on the composer's send button now, and the JS finds the
    running message through this attribute."""
    _add_connection(db)
    chat_id = make_chat()
    response = client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    assert "sse-connect" in response.text
    assert f"/api/chats/{chat_id}/messages/" in response.text


def test_the_streaming_bubble_renders_markdown_not_raw_tokens(
    client: TestClient, db, registered, make_chat
):
    """The body receives re-rendered Markdown, so formatting appears as the
    model writes rather than snapping in at the end."""
    _add_connection(db)
    chat_id = make_chat()
    response = client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    assert 'sse-swap="render"' in response.text
    assert 'hx-swap="innerHTML"' in response.text


def test_reasoning_starts_closed(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    response = client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    block = response.text[response.text.index("reasoning--live"):]
    assert not block[: block.index(">")].strip().endswith("open")


# --- Rewinding ---------------------------------------------------------------
def _exchange(client: TestClient, db, chat_id: str, text: str) -> Message:
    client.post(f"/api/chats/{chat_id}/messages", data={"content": text})
    assistant = db.scalars(
        select(Message).where(Message.role == "assistant").order_by(Message.created_at)
    ).all()[-1]
    assistant.content = f"reply to {text}"
    assistant.complete = True
    db.commit()
    return assistant


def test_editing_rewinds_and_discards_later_messages(
    client: TestClient, db, registered, make_chat
):
    _add_connection(db)
    chat_id = make_chat()
    _exchange(client, db, chat_id, "first")
    _exchange(client, db, chat_id, "second")
    assert db.scalar(select(func.count()).select_from(Message)) == 4

    first_user = db.scalars(
        select(Message).where(Message.role == "user").order_by(Message.created_at)
    ).first()
    client.post(
        f"/api/chats/{chat_id}/messages/{first_user.id}/edit",
        data={"content": "first, revised"},
    )

    # The edit happened in the request's session; this one still holds the old
    # instance in its identity map.
    db.expire_all()
    remaining = db.scalars(select(Message).order_by(Message.created_at)).all()
    assert [m.role for m in remaining] == ["user", "assistant"]
    assert remaining[0].content == "first, revised"
    # The fresh assistant row is incomplete, which is what restarts the stream.
    assert remaining[1].complete is False


def test_the_edit_form_says_how_much_will_be_lost(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    _exchange(client, db, chat_id, "first")
    _exchange(client, db, chat_id, "second")

    first_user = db.scalars(
        select(Message).where(Message.role == "user").order_by(Message.created_at)
    ).first()
    page = client.get(f"/api/chats/{chat_id}/messages/{first_user.id}/edit").text
    assert "3 messages after this one will be deleted" in page


def test_only_your_own_turns_can_be_edited(client: TestClient, db, registered, make_chat):
    """Rewriting what the model said would be inventing history."""
    _add_connection(db)
    chat_id = make_chat()
    assistant = _exchange(client, db, chat_id, "hello")
    assert client.get(
        f"/api/chats/{chat_id}/messages/{assistant.id}/edit"
    ).status_code == 404


def test_an_edit_cannot_empty_a_message(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    _exchange(client, db, chat_id, "hello")
    user_message = db.scalar(select(Message).where(Message.role == "user"))

    assert client.post(
        f"/api/chats/{chat_id}/messages/{user_message.id}/edit", data={"content": "   "}
    ).status_code == 400


def test_cancelling_an_edit_restores_the_bubble(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    _exchange(client, db, chat_id, "unchanged")
    user_message = db.scalar(select(Message).where(Message.role == "user"))

    page = client.get(f"/api/chats/{chat_id}/messages/{user_message.id}/cancel-edit").text
    assert "unchanged" in page
    assert "edit-form" not in page


# --- Background generation ---------------------------------------------------
def test_sending_launches_the_generation_immediately(
    client: TestClient, db, registered, make_chat
):
    """The reply is produced by a task, not by the browser watching it. That is
    what lets you navigate away without cutting it off."""
    from lembas.services import generation as generation_service

    _add_connection(db)
    chat_id = make_chat()
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))

    assert generation_service.get(message.id) is not None


def test_starting_a_chat_launches_the_generation(client: TestClient, db, registered):
    from lembas.services import generation as generation_service

    _add_connection(db)
    client.post("/api/chats/start", data={"content": "hi"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))
    assert generation_service.get(message.id) is not None


def test_asking_twice_does_not_start_a_second_generation(
    client: TestClient, db, registered, make_chat
):
    """A page load finding an unfinished reply must attach, not restart."""
    from lembas.services import generation as generation_service

    _add_connection(db)
    chat_id = make_chat()
    client.post(f"/api/chats/{chat_id}/messages", data={"content": "hi"})
    message = db.scalar(select(Message).where(Message.role == "assistant"))

    first = generation_service.get(message.id)
    assert generation_service.ensure(chat_id, message.id) is first


# --- Unread -------------------------------------------------------------------
def test_the_unread_poll_reports_dots(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    chat = db.get(Chat, chat_id)
    chat.unread = True
    db.commit()

    response = client.get("/api/chats/unread")
    assert f'id="unread-{chat_id}"' in response.text
    assert "hidden" not in response.text
    assert "lembas:unread" in response.headers.get("HX-Trigger", "")


def test_an_arrival_is_announced_once(client: TestClient, db, registered, make_chat):
    """Otherwise the same reply would toast every ten seconds forever."""
    _add_connection(db)
    chat_id = make_chat()
    chat = db.get(Chat, chat_id)
    chat.unread = True
    db.commit()

    assert "HX-Trigger" in client.get("/api/chats/unread").headers
    assert "HX-Trigger" not in client.get("/api/chats/unread").headers


def test_opening_a_chat_marks_it_read(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    chat = db.get(Chat, chat_id)
    chat.unread = True
    db.commit()

    client.get(f"/chat/{chat_id}")
    db.expire_all()
    assert db.get(Chat, chat_id).unread is False


def test_a_read_chat_reports_a_hidden_dot(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    response = client.get("/api/chats/unread")
    assert f'id="unread-{chat_id}"' in response.text
    assert "hidden" in response.text


def test_the_unread_poll_only_sees_your_own_chats(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    mine = make_chat()

    client.post("/auth/logout", follow_redirects=False)
    client.post(
        "/auth/register",
        data={"name": "Sam", "email": "sam@shire.test", "password": "potatoes-po-ta-toes"},
        follow_redirects=False,
    )
    assert mine not in client.get("/api/chats/unread").text


def test_the_sidebar_shows_an_unread_dot(client: TestClient, db, registered, make_chat):
    _add_connection(db)
    chat_id = make_chat()
    chat = db.get(Chat, chat_id)
    chat.unread = True
    db.commit()

    # Rendered on another page, so the dot is visible while looking elsewhere.
    page = client.get("/chat").text
    assert f'id="unread-{chat_id}" class="unread-dot"' in page
    assert 'hx-get="/api/chats/unread"' in page