Record streams the caller abandons

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.

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}"