pref: Improved _watch_disconnect CPU impact

src/mlpa/core/completions.py

@@ -114,6 +114,15 @@ async def stream_completion(
     async def _watch_disconnect() -> None:
         task = asyncio.current_task()
         assert task is not None
+        message = await request.receive()
+        if message["type"] == "http.disconnect":
+            disconnect_event.set()
+            return
+
+        log.warning(
+            "Unexpected ASGI message in disconnect watcher: {}", message["type"]
+        )
+
         while not await request.is_disconnected():
             # Starlette's AnyIO cancel scope can swallow task cancellation.
             # Honor the pending cancellation before starting another poll.

src/mlpa/run.py

@@ -11,6 +11,7 @@
 from fastapi.responses import StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
+from sentry_sdk.integrations.starlette import StarletteIntegration
 
 from mlpa.core.auth.authorize import authorize_chat_request, authorize_search_request
 from mlpa.core.classes import AuthorizedChatRequest, AuthorizedSearchRequest
@@ -121,6 +122,7 @@ def sentry_scrub_sensitive_fields(event, hint):
     before_send=sentry_scrub_sensitive_fields,
     dsn=env.SENTRY_DSN,
     send_default_pii=False,
+    integrations=[StarletteIntegration(middleware_spans=False)],
 )
 
 app = FastAPI(

src/tests/conftest.py

@@ -1,3 +1,4 @@
+import asyncio
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -14,10 +15,18 @@ def _reset_service_type_virtual_keys():
     env.__dict__.pop("service_type_virtual_keys", None)
 
 
+async def receive_until_cancelled():
+    """ASGI receive for a client that stays connected: blocks until cancelled."""
+    await asyncio.Event().wait()
+
+
 @pytest.fixture
 def mock_request():
     req = MagicMock()
     req.headers = {}
+    # The disconnect watcher awaits receive(). An AsyncMock would return a
+    # non-disconnect message and silently route tests through the polling fallback.
+    req.receive = receive_until_cancelled
     req.is_disconnected = AsyncMock(return_value=False)
     return req
 

src/tests/unit/test_completions.py

@@ -1951,14 +1951,21 @@ async def _aiter_bytes():
 
 async def test_response_aclose_called_once_on_disconnect(mocker, mock_request):
     role_chunk = b'data: {"choices":[{"delta":{"role":"assistant"}}]}\n\n'
+    first_chunk_forwarded = asyncio.Event()
 
     async def _aiter_bytes():
         yield role_chunk
+        first_chunk_forwarded.set()
         await asyncio.Event().wait()  # blocks: simulate waiting on next chunk
 
+    async def _receive():
+        # Client disconnects while the stream waits on the next upstream chunk.
+        await first_chunk_forwarded.wait()
+        return {"type": "http.disconnect"}
+
     capture = {}
     _patch_mock_stream_client(mocker, _aiter_bytes, capture)
-    mock_request.is_disconnected = AsyncMock(side_effect=[False, True])
+    mock_request.receive = _receive
 
     received = [c async for c in stream_completion(mock_request, SAMPLE_REQUEST)]
 
@@ -2117,7 +2124,7 @@ async def test_stream_cleanup_finishes_without_client_disconnect(
     )
     mocker.patch.object(env, "DISCONNECT_POLL_INTERVAL_SECONDS", 0.1)
     messages = asyncio.Queue()
-    poll_during_close = asyncio.Event()
+    watcher_in_receive = asyncio.Event()
     closing_upstream = False
     watch_task = None
     sent = []
@@ -2128,9 +2135,12 @@ async def test_stream_cleanup_finishes_without_client_disconnect(
 
     async def receive():
         nonlocal watch_task
-        watch_task = asyncio.current_task()
-        if closing_upstream:
-            poll_during_close.set()
+        task = asyncio.current_task()
+        # Under ASGI 2.3 Starlette's disconnect listener also calls receive();
+        # only count the watcher.
+        if task.get_coro().__name__ == "_watch_disconnect":
+            watch_task = task
+            watcher_in_receive.set()
         return await messages.get()
 
     async def upstream_chunks():
@@ -2140,7 +2150,9 @@ async def upstream_chunks():
     async def close_upstream():
         nonlocal closing_upstream
         closing_upstream = True
-        await poll_during_close.wait()
+        # Close the upstream only once the watcher is blocked in receive(),
+        # so cleanup cancels it there.
+        await watcher_in_receive.wait()
         if not cancel_during_poll:
             # Allow the poll's internal cancellation to finish first.
             await asyncio.sleep(0.005)

src/tests/unit/test_request_timings.py

@@ -6,9 +6,10 @@
 from mlpa.core.request_timings import RequestTimingsMiddleware, measure
 
 
-def make_request(value=None):
+def make_request(value=None, receive=None):
     headers = [] if value is None else [(b"debug-timing", value.encode())]
-    return Request({"type": "http", "method": "POST", "path": "/", "headers": headers})
+    scope = {"type": "http", "method": "POST", "path": "/", "headers": headers}
+    return Request(scope, receive) if receive else Request(scope)
 
 
 @pytest.mark.parametrize("value", [None, "", "false", "0", "1", "yes", "invalid"])

src/tests/unit/test_stream_timings.py

@@ -5,6 +5,7 @@
 
 from mlpa.core.completions import stream_completion
 from mlpa.core.request_timings import timing_snapshot
+from tests.conftest import receive_until_cancelled
 from tests.consts import SAMPLE_REQUEST
 from tests.unit.test_request_timings import make_request
 
@@ -15,8 +16,7 @@ async def test_stream_handshake_and_first_chunk_boundaries(mocker, metrics_spy,
     mocker.patch(
         "mlpa.core.request_timings.time.perf_counter", side_effect=lambda: now[0]
     )
-    request = make_request("true")
-    request.is_disconnected = mocker.AsyncMock(return_value=False)
+    request = make_request("true", receive=receive_until_cancelled)
     request.state.debug_timing_start = 1.0
     request.state.debug_spans = []
     closed = []