"""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 __enter__(self) -> _FakeSession: return self def __exit__(self, *args: Any) -> None: pass 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). SEC-14-04: the chat endpoint uses short-lived sessions via ``SessionLocal()`` — we monkeypatch ``chat_api.SessionLocal`` to return a fake session instead of overriding ``get_db``. """ monkeypatch.setattr(chat_api, "db_available", lambda: True) session = _FakeSession() monkeypatch.setattr(chat_api, "SessionLocal", 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 )