Fix generation request cleanup

gradio_ui.py

@@ -33,7 +33,15 @@ async def events(response):
                 raw, pending = pending.split("\n\n", 1)
                 for line in raw.splitlines():
                     if line.startswith("data: "):
-                        yield json.loads(line[6:])
+                        event = json.loads(line[6:])
+                        if event.get("type") in {"done", "error"}:
+                            # Gradio can pause this adapter while rendering its
+                            # final update. Release backend slots before the UI
+                            # lets the user start another take.
+                            await response.body_iterator.aclose()
+                            yield event
+                            return
+                        yield event
 
 
 class Studio:

server.py

@@ -682,8 +682,8 @@ class Gate:
 class ClientLimits:
     """Per-account concurrency and burst limits.
 
-    A slot expires after SLOT_TTL_S even if never released: a stream whose client
-    left before its body started never runs its cleanup.
+    Expiration is a final recovery guard. Streams reserve their slot only once
+    consumption starts, and release it when their iterator closes.
     """
 
     SLOT_TTL_S = 300
@@ -692,20 +692,32 @@ class ClientLimits:
         self.active: dict[str, list[float]] = collections.defaultdict(list)
         self.recent: dict[str, collections.deque[float]] = collections.defaultdict(collections.deque)
 
-    def admit(self, client: str) -> float | str:
-        """A slot token, or the reason the request is refused."""
+    def check(self, client: str) -> str | None:
+        """Check limits without reserving a slot or recording a new request."""
         now = time.time()
         recent = self.recent[client]
         while recent and now - recent[0] > RATE_WINDOW_S:
             recent.popleft()
-        slots = self.active[client] = [t for t in self.active[client] if now - t < self.SLOT_TTL_S]
+        slots = [t for t in self.active.get(client, []) if now - t < self.SLOT_TTL_S]
+        if slots:
+            self.active[client] = slots
+        else:
+            self.active.pop(client, None)
         if len(slots) >= PER_CLIENT_ACTIVE:
-            return "One take at a time, please. Wait for the current one to finish."
+            return f"You already have {PER_CLIENT_ACTIVE} takes in progress. Wait for one to finish."
         if len(recent) >= PER_CLIENT_REQUESTS:
             wait = int(RATE_WINDOW_S - (now - recent[0])) + 1
             return f"That's a lot of takes! Give the GPU a breather and try again in {wait}s."
-        recent.append(now)
-        slots.append(now)
+        return None
+
+    def admit(self, client: str) -> float | str:
+        """Reserve atomically when a response actually begins streaming."""
+        refused = self.check(client)
+        if refused:
+            return refused
+        now = time.time()
+        self.recent[client].append(now)
+        self.active[client].append(now)
         return now
 
     def release(self, client: str, slot: float) -> None:
@@ -1138,14 +1150,18 @@ async def generate(body: GenerateBody, identity: Identity = Depends(identify)) -
     if len(gate.waiting) >= MAX_QUEUE:
         raise HTTPException(503, "The stage is packed right now. Try again in a minute.")
     await asyncio.to_thread(check_quota, identity)
-    slot = limits.admit(identity.key)
-    if isinstance(slot, str):
-        raise HTTPException(429, slot)
+    refused = limits.check(identity.key)
+    if refused:
+        raise HTTPException(429, refused)
 
     payload = engine_request(choice, turn_text, history, body)
     take_id = secrets.token_urlsafe(9)
 
     async def stream():
+        slot = limits.admit(identity.key)
+        if isinstance(slot, str):
+            yield sse({"type": "error", "message": slot})
+            return
         ticket = gate.enqueue()
         queued_at = time.perf_counter()
         try:
@@ -1162,10 +1178,11 @@ async def generate(body: GenerateBody, identity: Identity = Depends(identify)) -
                 try:
                     while not waking.done():
                         if not (await asyncio.wait({waking}, timeout=KEEPALIVE_S))[0]:
-                            yield sse({"type": "retrying", "message": "Waiting for the H200 endpoint to wake up…"})
+                            elapsed = round(time.perf_counter() - admitted)
+                            yield sse({"type": "retrying", "message": f"Connecting to the H200 endpoint… ({elapsed} s)"})
                     waking.result()
                 except asyncio.TimeoutError:
-                    yield sse({"type": "error", "message": "The H200 endpoint is still starting. Try again in a moment."})
+                    yield sse({"type": "error", "message": "Could not connect to the H200 endpoint. Try again in a moment."})
                     return
                 finally:
                     if not waking.done():

tests/test_request_cleanup.py

@@ -0,0 +1,67 @@
+"""Request leases must not outlive a cancelled or terminal Gradio response."""
+from contextlib import aclosing, AsyncExitStack
+import json
+
+import pytest
+
+from test_gradio_ui import anyio_backend, studio, browser_request, body, EngineClient, EngineResponse
+import server
+from accounts import Identity, LOCAL
+
+
+def identity():
+    return Identity(keys=("test",), name="test", entitlement=LOCAL, signed_in=False)
+
+
+@pytest.mark.anyio
+async def test_unstarted_response_does_not_reserve_a_slot(studio):
+    response = await server.generate(body(studio), identity())
+    try:
+        assert not server.limits.active
+        assert server.gate.active == 0
+    finally:
+        await response.body_iterator.aclose()
+
+
+@pytest.mark.anyio
+async def test_payload_failure_cannot_leak_a_slot(studio, monkeypatch):
+    def invalid_payload(*args):
+        raise ValueError("Invalid reference")
+    monkeypatch.setattr(server, "engine_request", invalid_payload)
+    with pytest.raises(ValueError, match="Invalid reference"):
+        await server.generate(body(studio), identity())
+    assert not server.limits.active
+
+
+@pytest.mark.anyio
+async def test_prepared_responses_still_enforce_concurrency_when_started(studio, monkeypatch):
+    monkeypatch.setattr(server, "PER_CLIENT_ACTIVE", 2)
+    monkeypatch.setattr(server, "http", EngineClient(EngineResponse()))
+    responses = [await server.generate(body(studio), identity()) for _ in range(3)]
+    async with AsyncExitStack() as stack:
+        streams = [await stack.enter_async_context(aclosing(r.body_iterator)) for r in responses]
+        await anext(streams[0])
+        await anext(streams[1])
+        refused = json.loads((await anext(streams[2])).removeprefix("data: "))
+        assert refused["type"] == "error" and "in progress" in refused["message"]
+        assert len(server.limits.active[identity().key]) == 2
+        assert server.gate.active == 2
+    assert not server.limits.active and server.gate.active == 0
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("error", [False, True])
+async def test_slot_is_free_before_terminal_event_reaches_browser(studio, browser_request, monkeypatch, error):
+    response = EngineResponse(error=error)
+    monkeypatch.setattr(server, "http", EngineClient(response))
+    async with aclosing(studio.stage_run(body(studio), None, [], browser_request)) as stream:
+        async for result in stream:
+            if result[0]["type"] in {"done", "error"}:
+                # The browser can click Perform as soon as this update arrives,
+                # while Gradio is still processing the outer iterator.
+                assert not server.limits.active
+                assert server.gate.active == 0
+                assert response.closed
+                break
+        else:
+            pytest.fail("No terminal event")