feat(chat): stop an in-flight answer — Send becomes Stop, the partial is kept and persisted, the model stream is torn down
This commit is contained in:
@@ -0,0 +1,419 @@
|
||||
"""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 gc
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, MutableMapping
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from app.api import chat as chat_api
|
||||
from app.config import 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,
|
||||
)
|
||||
|
||||
|
||||
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 _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"),
|
||||
],
|
||||
"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
|
||||
)
|
||||
Reference in New Issue
Block a user