fix(sse): stop connection slots leaking when a client disconnects early

app/api/events.py

@@ -38,6 +38,45 @@ def release_sse_slot() -> None:
     _sse_active -= 1
 
 
+class SseSlot:
+    """One held connection slot, released exactly once.
+
+    Claiming has to happen in the handler so that hitting the cap can still be
+    answered with a 503 -- once the generator is running the response headers
+    have gone out and there is no status code left to send.
+
+    That is what leaked slots: release lives in the stream's `finally`, and an
+    async generator that is never started never runs its `finally`. If the
+    client disconnects before the body begins, StreamingResponse raises inside
+    `stream_response` on its first `send()` -- before `__anext__` is ever
+    called -- so the generator body never executes and the slot was held
+    forever. 200 of those and every progress stream 503s with nothing actually
+    connected, until the process restarts (#513).
+
+    The stream releases on its way out as before; `__del__` is the backstop for
+    the never-started case, where collecting the generator collects the closure
+    holding this. Release is idempotent so the two cannot double-count.
+    """
+
+    __slots__ = ("_held",)
+
+    def __init__(self) -> None:
+        # Set first: claim_sse_slot raises at the cap, and __del__ still runs on
+        # a half-built object. Without this it would raise AttributeError from
+        # __del__ instead of releasing nothing.
+        self._held = False
+        claim_sse_slot()  # may raise 503; nothing is held if it does
+        self._held = True
+
+    def release(self) -> None:
+        if self._held:
+            self._held = False
+            release_sse_slot()
+
+    def __del__(self) -> None:
+        self.release()
+
+
 @router.get("/jobs/{job_id}/events")
 async def job_events(job_id: str) -> StreamingResponse:
     """Server-Sent Events stream of job state updates. Closes when the job
@@ -47,7 +86,7 @@ async def job_events(job_id: str) -> StreamingResponse:
     job = registry_get(job_id)
     if job is None:
         raise HTTPException(status_code=404, detail="job not found")
-    claim_sse_slot()
+    slot = SseSlot()
 
     async def stream() -> AsyncIterator[str]:
         try:
@@ -82,7 +121,7 @@ async def stream() -> AsyncIterator[str]:
                     keepalive_at = 0
                 await asyncio.sleep(0.2)
         finally:
-            release_sse_slot()
+            slot.release()
 
     return StreamingResponse(
         stream(),

app/api/queue.py

@@ -22,7 +22,7 @@
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
-from app.api.events import _MAX_SSE_SECONDS, claim_sse_slot, release_sse_slot
+from app.api.events import _MAX_SSE_SECONDS, SseSlot
 from app.core.config import JOB_ID_RE, JOBS_DIR, MAX_PENDING_UPLOAD_JOBS, MAX_PENDING_URL_JOBS
 from app.core.registry import get as registry_get
 from app.core.registry import pending_count as registry_pending_count
@@ -130,7 +130,7 @@ async def queue_events() -> StreamingResponse:
     outlives any individual job and is expected to stay open for the session,
     so only the 4 h ceiling ends it.
     """
-    claim_sse_slot()
+    slot = SseSlot()
 
     async def stream() -> AsyncIterator[str]:
         try:
@@ -156,7 +156,7 @@ async def stream() -> AsyncIterator[str]:
                     keepalive_at = 0
                 await asyncio.sleep(0.25)
         finally:
-            release_sse_slot()
+            slot.release()
 
     return StreamingResponse(
         stream(),

tests/test_sse_slot_budget.py

@@ -0,0 +1,92 @@
+"""The shared SSE connection budget must not leak (#513).
+
+claim_sse_slot() runs in the handler so hitting the cap can still answer 503.
+release lives in the stream's finally -- and an async generator that is never
+started never runs its finally. A client that disconnects before the response
+body begins therefore held a slot forever: 200 of those and every progress
+stream 503s with nothing actually connected, until the process restarts.
+"""
+
+from __future__ import annotations
+
+import gc
+
+import pytest
+from fastapi import HTTPException
+
+from app.api import events as _events
+
+
+@pytest.fixture(autouse=True)
+def _zero_budget():
+    _events._sse_active = 0
+    yield
+    _events._sse_active = 0
+
+
+def test_a_slot_is_held_then_released():
+    slot = _events.SseSlot()
+    assert _events._sse_active == 1
+
+    slot.release()
+    assert _events._sse_active == 0
+
+
+def test_release_is_idempotent():
+    # The stream's finally and the __del__ backstop can both fire; counting
+    # twice would free a slot that is still in use.
+    slot = _events.SseSlot()
+    slot.release()
+    slot.release()
+
+    assert _events._sse_active == 0
+
+
+def test_a_slot_dropped_without_release_is_reclaimed():
+    # The leak: the generator is created, so the slot is claimed, but never
+    # iterated, so its finally never runs. Collecting it must free the slot.
+    def _never_started():
+        slot = _events.SseSlot()
+
+        async def stream():
+            try:
+                yield "data: x\n\n"
+            finally:
+                slot.release()
+
+        return stream()  # created, never iterated
+
+    gen = _never_started()
+    assert _events._sse_active == 1
+
+    del gen
+    gc.collect()
+
+    assert _events._sse_active == 0, "a slot leaked for the life of the process"
+
+
+def test_the_cap_still_answers_503():
+    held = [_events.SseSlot() for _ in range(_events._MAX_SSE_CONNECTIONS)]
+    assert _events._sse_active == _events._MAX_SSE_CONNECTIONS
+
+    with pytest.raises(HTTPException) as excinfo:
+        _events.SseSlot()
+    assert excinfo.value.status_code == 503
+
+    for slot in held:
+        slot.release()
+    assert _events._sse_active == 0
+
+
+def test_a_refused_claim_holds_nothing():
+    # If __init__ raises, no slot was taken -- so the failed attempt must not
+    # decrement on collection either.
+    held = [_events.SseSlot() for _ in range(_events._MAX_SSE_CONNECTIONS)]
+    with pytest.raises(HTTPException):
+        _events.SseSlot()
+
+    gc.collect()
+    assert _events._sse_active == _events._MAX_SSE_CONNECTIONS
+
+    for slot in held:
+        slot.release()