LLeMbas/tests/test_chat.py

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

from __future__ import annotations

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"


# --- 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):
    _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_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()

    contents = [m["content"] for m in chat_service.build_request(db, chat)["messages"]]
    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", "content": "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