diff --git a/src/api/answer.py b/src/api/answer.py index 3cfd753..2c12a5f 100644 --- a/src/api/answer.py +++ b/src/api/answer.py @@ -123,6 +123,7 @@ async def stream() -> AsyncIterator[str]: # separate events, so they come out here -- across fragment boundaries, # because one anchor arrives as twenty-odd fragments. stripper = AnchorStripper() + tokens_sent = 0 try: async with asyncio.timeout(ANSWER_TIMEOUT_SECONDS): async for event in graph.astream_answer( @@ -140,6 +141,7 @@ async def stream() -> AsyncIterator[str]: if event.kind == "token": text = stripper.feed(event.text) if text: + tokens_sent += 1 yield _sse("token", {"text": text}) elif event.kind == "citation": yield _sse( @@ -151,6 +153,33 @@ async def stream() -> AsyncIterator[str]: ) elif event.kind == "done": state = event.state or "failed" + except (asyncio.CancelledError, GeneratorExit): + # The caller hung up. Both forms are caught because a hang-up + # arrives as either, depending on who notices first: Starlette + # cancelling the task raises CancelledError, while closing the + # generator raises GeneratorExit. A test that only closed the + # generator passed against a handler catching only CancelledError, + # which is how this was found. + # + # Either way the graph is cancelled and the model call stops -- that + # part already worked and is measured. What was missing was a record. + # + # Worth a log line because an abandoned stream is indistinguishable + # from a healthy one in every other signal: the request 200s, tokens + # flow, and then nothing. The website found a bug on their side where + # a keystroke unmounted the panel mid-answer, and from here it would + # have looked like ordinary traffic. A rate of these is the symptom + # of a caller that starts answers it does not want. + # + # Re-raised, never swallowed: cancellation is not an error to report + # to a caller who has already gone, and suppressing it would leave + # the task pretending to still be running. + logger.info( + "answer abandoned by the caller after %.1fs and %d token events", + time.monotonic() - started, + tokens_sent, + ) + raise except TimeoutError: # Distinct from the crash below so an operator can tell a slow # upstream from a broken one. diff --git a/tests/api/test_answer_endpoint.py b/tests/api/test_answer_endpoint.py index 028cb97..5270b08 100644 --- a/tests/api/test_answer_endpoint.py +++ b/tests/api/test_answer_endpoint.py @@ -13,8 +13,8 @@ import asyncio import json import time -from collections.abc import AsyncIterator -from typing import Any +from collections.abc import AsyncGenerator, AsyncIterator +from typing import Any, cast import jwt import pytest @@ -24,7 +24,8 @@ from fastapi.testclient import TestClient from agent.graph import AnswerEvent -from api.answer import router +from api.answer import AnswerRequest, router +from api.answer import answer as answer_handler from util.caller_token import DEFAULT_AUDIENCE from util.rate_limit import SlidingWindowLimiter @@ -424,3 +425,66 @@ async def astream_answer( ) assert seen["enable_postprocess"] is False + + +def test_an_abandoned_stream_is_recorded_and_not_swallowed( + keys: tuple[str, str], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """A caller that hangs up mid-answer should leave a trace. + + Cancellation already stopped the work -- that was measured against a live + server, and it is why the website's proxy does not need to cancel upstream. + What was missing was any record: an abandoned stream looks exactly like a + healthy one in every other signal, because the request 200s, tokens flow, and + then nothing more happens. + + It matters because the website found a bug where a keystroke unmounted their + panel mid-answer and left our stream running. From here that was ordinary + traffic. A rate of these is the symptom. + + Driven through the response's iterator rather than a client, because the + point is to close it mid-stream, which a TestClient will not do. + """ + from types import SimpleNamespace + + private, public = keys + + class _SlowGraph(_StubGraph): + async def astream_answer( + self, *_a: Any, **_k: Any + ) -> AsyncIterator[AnswerEvent]: + for index in range(50): + await asyncio.sleep(0.01) + yield AnswerEvent(kind="token", text=f"t{index} ") + + monkeypatch.setattr("api.answer.get_graph", lambda: _SlowGraph()) + request = SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(caller_token_key=public, release=97)) + ) + + async def drive() -> None: + response = await answer_handler( + request, # type: ignore[arg-type] + AnswerRequest(question="what is CDK5", caller_token=_token(private)), + ) + # body_iterator is typed as a bare AsyncIterable, which has no aclose; + # the object Starlette puts there is an async generator, and closing it + # is the whole point of this test. + iterator = cast("AsyncGenerator[str, None]", response.body_iterator) + seen = 0 + async for _chunk in iterator: + seen += 1 + if seen == 4: + break + # Closing is what a hung-up caller causes Starlette to do. + await iterator.aclose() + + with caplog.at_level("INFO", logger="api.answer"): + asyncio.run(drive()) + + messages = [record.getMessage() for record in caplog.records] + assert any( + "abandoned by the caller" in message for message in messages + ), f"no record of the abandoned stream; logged: {messages}"