brain-of-reese/tests/unit/test_chat_cancel.py

"""Unit: client-disconnect teardown of a chat turn (phase 48, task 01).

Drives ``POST /api/chat`` through the real ASGI app — a real
``LLMClient`` with a slow, recording fake SDK stream behind it, plus the
fake DB session / retriever pattern from ``tests/unit/test_chat_gate.py``
— with an ASGI-level client disconnect: after a few SSE frames the
``receive`` channel starts returning ``http.disconnect``, the ASGI 2.0
contract this Starlette's ``StreamingResponse`` listens for (its task
group cancels the body task, and the abandoned body generator chain is
finalized by the loop's PEP 525 asyncgen hooks). The ``TestClient``
transport buffers the whole body and cannot drop a connection
mid-stream, so the disconnect is emulated at the ASGI boundary —
exactly where a real server hands it over.

Owner-locked contract (2026-08-29): on a cancelled turn the model's
stream is closed promptly, one ``chat: turn cancelled`` log line is
written, **no** ``query_log`` row exists, and no ``done``/``error``
frame follows the disconnect — while completed turns and the
mid-stream ``LLMError`` path settle exactly as before (``done`` frame +
query_log row; structured ``error`` frame, not logged as cancelled).
"""
from __future__ import annotations

import asyncio
import base64
import gc
import json
import logging
import time
import uuid
from collections.abc import Callable, Iterator, MutableMapping
from datetime import UTC, datetime
from types import SimpleNamespace
from typing import Any

import pytest
from itsdangerous import TimestampSigner

from app.api import chat as chat_api
from app.config import Settings, get_settings
from app.main import app as fastapi_app
from app.models import Document, KbOverview, QueryLog
from app.rag.llm import LLMClient
from app.rag.retriever import RetrievedChunk


def _doc(title: str, content: str) -> Document:
    return Document(
        id=uuid.uuid4(),
        source="Homelab",
        path=f"{title.lower().replace(' ', '-')}.md",
        full_path="/tmp/doc.md",
        title=title,
        content=content,
        content_hash="0" * 64,
        # Phase 106, D5: the HIGH block formats the row's created_at
        # UTC date part — the detached fixture carries it (the NOT NULL
        # DB column guarantees it for real rows).
        created_at=datetime(2024, 6, 15, 12, 0, 0, tzinfo=UTC),
    )


def _chunk(doc: Document, cosine: float, fts_hit: bool = False) -> RetrievedChunk:
    return RetrievedChunk(
        chunk_id=uuid.uuid4(),
        position=0,
        content=doc.content[:32],
        score=cosine,
        document=doc,
        cosine=cosine,
        fts_hit=fts_hit,
    )


def _fake_retriever(chunks: list[RetrievedChunk]) -> Any:
    def retrieve(_db: Any, _question: str, _vec: list[float]) -> list[RetrievedChunk]:
        return chunks

    return retrieve


# ---------- fakes: slow SDK stream behind a real LLMClient ----------


def _sse_chunk(text: str) -> SimpleNamespace:
    """One fake ChatCompletionChunk (``choices[].delta.content`` shape)."""
    return SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=text))])


class _SlowStream:
    """A fake aipi SSE stream (the openai SDK ``AsyncStream`` shape):
    yields *chunks* with a small sleep between them (so a disconnect can
    land mid-iteration) and records ``close()`` calls — the SDK stream's
    deterministic teardown. ``fail_after`` makes ``__anext__`` raise a
    transport error after that many chunks (the mid-stream
    ``LLMError`` path)."""

    def __init__(
        self, chunks: list, fail_after: int | None = None, delay: float = 0.005
    ) -> None:
        self._chunks = list(chunks)
        self._fail_after = fail_after
        self._delay = delay
        self._i = 0
        self.closed = False

    def __aiter__(self) -> _SlowStream:
        return self

    async def __anext__(self) -> SimpleNamespace:
        self._i += 1
        if self._fail_after is not None and self._i > self._fail_after:
            raise ConnectionError("simulated mid-stream drop")
        if self._i > len(self._chunks):
            raise StopAsyncIteration
        await asyncio.sleep(self._delay)
        return self._chunks[self._i - 1]

    async def close(self) -> None:
        self.closed = True


class _FakeCompletions:
    def __init__(self, stream: _SlowStream) -> None:
        self._stream = stream
        self.kwargs: dict | None = None

    async def create(self, **kwargs: Any):
        self.kwargs = kwargs
        assert kwargs.get("stream") is True
        return self._stream


def _make_llm(monkeypatch: pytest.MonkeyPatch, stream: _SlowStream) -> LLMClient:
    """A real ``LLMClient`` (so the production ``chat_stream`` teardown
    runs) with the fake SDK stream behind it and a deterministic
    ``embed_one`` (no embeddings HTTP)."""
    llm = LLMClient(Settings(_env_file=None))  # pyright: ignore[reportCallIssue]
    llm._client = SimpleNamespace(  # pyright: ignore[reportAttributeAccessIssue]
        chat=SimpleNamespace(completions=_FakeCompletions(stream))
    )

    async def _embed_one(self: Any, _text: str) -> list[float]:
        return [0.0] * 768

    monkeypatch.setattr(LLMClient, "embed_one", _embed_one)
    return llm


# ---------- fake DB session (test_chat_gate.py pattern) ----------


class _FakeSteeringResult:
    def all(self) -> list[Any]:
        return []


class _FakeSession:
    """Records the QueryLog rows it is given; no steering notes, no
    stored KB overview."""

    def __init__(self) -> None:
        self.added: list[Any] = []
        self.commits = 0

    def add(self, obj: Any) -> None:
        self.added.append(obj)

    def commit(self) -> None:
        self.commits += 1

    def scalars(self, _stmt: Any) -> _FakeSteeringResult:
        return _FakeSteeringResult()

    def get(self, model: Any, pk: Any) -> Any:
        if model is KbOverview:
            return KbOverview(id=1, content="")
        return None


@pytest.fixture()
def env(monkeypatch: pytest.MonkeyPatch) -> Iterator[_FakeSession]:
    """``POST /api/chat`` with the DB session, retriever settings, and
    availability faked (the gate tests' wiring)."""
    monkeypatch.setattr(chat_api, "db_available", lambda: True)
    session = _FakeSession()
    monkeypatch.setitem(fastapi_app.dependency_overrides, chat_api.get_db, lambda: session)
    # A stable gate threshold, independent of the production default.
    monkeypatch.setattr(
        chat_api,
        "get_settings",
        lambda: Settings(_env_file=None, relevance_threshold=0.30),  # pyright: ignore[reportCallIssue]
    )
    yield session
    fastapi_app.dependency_overrides.clear()


def _install_llm(monkeypatch: pytest.MonkeyPatch, llm: LLMClient) -> None:
    monkeypatch.setitem(fastapi_app.dependency_overrides, chat_api.get_llm, lambda: llm)


# ---------- the ASGI driver (client disconnect at the ASGI boundary) ----------


def _admin_cookie_header() -> tuple[bytes, bytes]:
    """A valid signed ``bor_session`` cookie carrying the admin session.

    Phase 79 (task 03): ``POST /api/chat`` is user-gated, and the raw
    ASGI scope below carries no browser — so it presents the same signed
    cookie ``SessionMiddleware`` would have emitted after
    ``POST /api/login`` (the admin session short-circuits
    ``require_user``; the anonymous 401 contract is pinned in
    ``test_auth_api.py``)."""
    settings = get_settings()
    data = base64.b64encode(json.dumps({"admin": True}).encode("utf-8"))
    signed = TimestampSigner(settings.session_secret).sign(data)
    return b"cookie", f"{settings.session_cookie}={signed.decode('ascii')}".encode("ascii")


def _scope() -> dict[str, Any]:
    return {
        "type": "http",
        "asgi": {"version": "3.0"},
        "http_version": "1.1",
        "method": "POST",
        "path": "/api/chat",
        "raw_path": b"/api/chat",
        "root_path": "",
        "scheme": "http",
        "query_string": b"",
        "headers": [
            (b"host", b"testserver"),
            (b"content-type", b"application/json"),
            _admin_cookie_header(),  # phase 79: the signed-in admin
        ],
        "client": ("testclient", 50000),
        "server": ("testserver", 80),
        "state": {},
    }


async def _drive(
    question: str,
    stop_after: int | None,
    settle: Callable[[], bool] | None = None,
) -> list[bytes]:
    """Run one ``POST /api/chat`` through the real ASGI app.

    ``stop_after=None`` lets the turn complete; otherwise the client
    "disconnects" after ``stop_after`` body chunks (SSE frames) have
    been written — the ``receive`` channel starts returning
    ``http.disconnect`` (the ASGI 2.0 contract; no ``spec_version`` in
    the scope, so ``StreamingResponse`` runs its task-group
    listen-for-disconnect path). When *settle* is given, spins the loop
    until it holds (the PEP 525 asyncgen finalizers run the abandoned
    generator chain's ``finally`` blocks a few loop turns after their
    frames are released) or a 5 s timeout runs out. Returns the body
    chunks written.
    """
    body = json.dumps({"message": question}).encode()
    chunks: list[bytes] = []
    disconnect = asyncio.Event()
    request_done = False

    async def receive() -> dict[str, Any]:
        nonlocal request_done
        if not request_done:
            request_done = True
            return {"type": "http.request", "body": body, "more_body": False}
        await disconnect.wait()
        return {"type": "http.disconnect"}

    async def send(message: MutableMapping[str, Any]) -> None:
        if message["type"] != "http.response.body":
            return
        chunk = message.get("body", b"")
        if chunk:
            chunks.append(chunk)
            if stop_after is not None and len(chunks) >= stop_after:
                disconnect.set()  # the client goes away

    await fastapi_app(_scope(), receive, send)
    if settle is not None:
        gc.collect()  # release any cycle-held frames up front
        deadline = time.monotonic() + 5.0
        while not settle():
            if time.monotonic() >= deadline:
                break
            await asyncio.sleep(0.01)
    return chunks


def _frames(chunks: list[bytes]) -> list[dict[str, Any]]:
    """Parse the SSE frames out of the written body chunks."""
    frames: list[dict[str, Any]] = []
    for raw in chunks:
        for frame in raw.decode("utf-8").split("\n\n"):
            frame = frame.strip()
            if frame.startswith("data:"):
                frames.append(json.loads(frame.removeprefix("data:").strip()))
    return frames


# ---------- cancelled turns (the phase-48 contract) ----------


def test_cancelled_grounded_turn_closes_stream_logs_cancel_and_skips_query_log(
    env: _FakeSession,
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Grounded turn (the agent loop): a mid-stream disconnect closes
    the model's stream, logs one cancel line, writes no query_log row,
    and emits no done/error frame after the disconnect."""
    stream = _SlowStream([_sse_chunk(f"word{i} ") for i in range(60)])
    llm = _make_llm(monkeypatch, stream)
    _install_llm(monkeypatch, llm)
    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_SENT")
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(doc, 0.90)]))

    with caplog.at_level(logging.INFO, logger="app.chat"):
        chunks = asyncio.run(
            _drive(
                "How is my Kubernetes cluster set up?",
                stop_after=2,
                settle=lambda: stream.closed,
            )
        )

    # The model's stream was closed promptly on abandon.
    assert stream.closed
    # The cancel log line — exactly once, with the question.
    cancel_lines = [
        r.getMessage() for r in caplog.records if "turn cancelled" in r.getMessage()
    ]
    assert len(cancel_lines) == 1
    assert "How is my Kubernetes cluster set up?" in cancel_lines[0]
    assert "total_ms=" in cancel_lines[0]
    # No durable record for a cancelled turn.
    assert env.added == []
    # Frames: the streamed deltas only — no done, no error, after the
    # disconnect (a third delta may race the teardown; all deltas).
    frames = _frames(chunks)
    assert len(frames) >= 2
    assert all(f["type"] == "delta" for f in frames)


def test_cancelled_deflected_turn_closes_stream(
    env: _FakeSession,
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """Deflected turn (the direct ``chat_stream`` path — A8): the same
    teardown contract holds without the agent loop."""
    stream = _SlowStream([_sse_chunk(f"word{i} ") for i in range(60)])
    llm = _make_llm(monkeypatch, stream)
    _install_llm(monkeypatch, llm)
    doc = _doc("Deploying a New Service", "DOC_CONTENT_NEVER_SENT")
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(doc, 0.10)]))

    with caplog.at_level(logging.WARNING, logger="app.chat"):
        chunks = asyncio.run(
            _drive(
                "How do I bake sourdough bread?",
                stop_after=2,
                settle=lambda: stream.closed,
            )
        )

    assert stream.closed
    cancel_lines = [
        r.getMessage() for r in caplog.records if "turn cancelled" in r.getMessage()
    ]
    assert len(cancel_lines) == 1
    assert env.added == []
    frames = _frames(chunks)
    assert len(frames) >= 2
    assert all(f["type"] == "delta" for f in frames)


# ---------- regressions: settled turns behave exactly as before ----------


def test_completed_turn_still_emits_done_and_writes_query_log(
    env: _FakeSession,
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """A completed turn: the ``done`` frame, the query_log row, and the
    per-turn log line — and NO cancel line (it settled)."""
    stream = _SlowStream([_sse_chunk(f"word{i} ") for i in range(3)], delay=0.001)
    llm = _make_llm(monkeypatch, stream)
    _install_llm(monkeypatch, llm)
    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_SENT")
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(doc, 0.90)]))

    with caplog.at_level(logging.INFO, logger="app.chat"):
        chunks = asyncio.run(_drive("How is my Kubernetes cluster set up?", None))

    frames = _frames(chunks)
    assert [f["type"] for f in frames] == ["delta", "delta", "delta", "done"]
    assert frames[-1]["deflected"] is False
    assert frames[-1]["sources"][0]["title"] == "Kubernetes Homelab Cluster"
    (row,) = env.added
    assert isinstance(row, QueryLog)
    assert row.question == "How is my Kubernetes cluster set up?"
    assert env.commits == 1
    # The per-turn line still goes out; no cancel line for a settled turn.
    assert any("question=" in r.getMessage() for r in caplog.records)
    assert not any("turn cancelled" in r.getMessage() for r in caplog.records)


def test_mid_stream_llm_error_settles_not_cancelled(
    env: _FakeSession,
    monkeypatch: pytest.MonkeyPatch,
    caplog: pytest.LogCaptureFixture,
) -> None:
    """The mid-stream ``LLMError`` path (the fake stream drops after two
    pieces): the structured ``error`` frame is emitted, the turn is NOT
    logged as cancelled (it settled), the stream is closed on the
    exception path, and — as before — no query_log row is written."""
    stream = _SlowStream([_sse_chunk(f"word{i} ") for i in range(60)], fail_after=2)
    llm = _make_llm(monkeypatch, stream)
    _install_llm(monkeypatch, llm)
    doc = _doc("Kubernetes Homelab Cluster", "TALOS_DOC_SENT")
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(doc, 0.90)]))

    with caplog.at_level(logging.INFO, logger="app.chat"):
        chunks = asyncio.run(_drive("How is my Kubernetes cluster set up?", None))

    frames = _frames(chunks)
    assert [f["type"] for f in frames] == ["delta", "delta", "error"]
    assert frames[-1]["detail"] == "The chat model dropped the connection — try again?"
    # The exception path closes the model's stream too (synchronously —
    # no settle needed).
    assert stream.closed
    assert env.added == []
    # Settled: no cancel line (the LLM stream failure IS logged, though).
    assert not any("turn cancelled" in r.getMessage() for r in caplog.records)
    assert any(
        "LLM stream failed" in r.getMessage() for r in caplog.records
    )