fix: close upstream LLM stream on client disconnect (port BerriAI#30245) (#184)

On client disconnect mid-stream, Starlette abandons the response body iterator without calling aclose(), leaving the upstream vLLM connection open and the GPU slot held until GC (zombie streams -> pool/slot starvation -> 540s timeouts under load).

Wrap the streaming response so its __call__ finally force-closes both the body iterator and the upstream generator, triggering the existing response.aclose() cleanup; add BaseModelResponseIterator.aclose()/http_response so the close reaches the live connection. Both proxy streaming paths (/chat/completions and /v1/messages) are covered via create_response.

Ports BerriAI#30245. Verified: 92 tests pass; ruff clean; black-clean on added lines.

Co-authored-by: Claude <noreply@anthropic.com>

Author: prakhar-prakash-juspay

litellm/llms/base_llm/base_model_iterator.py

@@ -1,8 +1,11 @@
 import json
 from abc import abstractmethod
-from typing import List, Optional, Union, cast
+from typing import TYPE_CHECKING, List, Optional, Union, cast
 
 import litellm
+
+if TYPE_CHECKING:
+    import httpx
 from litellm.types.utils import (
     Choices,
     Delta,
@@ -64,6 +67,18 @@ def __init__(
         self.streaming_response = streaming_response
         self.response_iterator = self.streaming_response
         self.json_mode = json_mode
+        self.http_response: Optional["httpx.Response"] = None
+
+    async def aclose(self) -> None:
+        """Close the upstream HTTP response so the provider connection is
+        released (and a backend like vLLM aborts generation) when the stream
+        is abandoned before its natural end.
+
+        ``streaming_response`` is usually a bare ``aiter_lines()`` generator
+        that holds no reference to the response, so the handler that owns the
+        response attaches it here after construction."""
+        if self.http_response is not None:
+            await self.http_response.aclose()
 
     def chunk_parser(
         self, chunk: dict

litellm/llms/custom_httpx/llm_http_handler.py

@@ -31,7 +31,10 @@
 from litellm.llms.base_llm.audio_transcription.transformation import (
     BaseAudioTranscriptionConfig,
 )
-from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
+from litellm.llms.base_llm.base_model_iterator import (
+    BaseModelResponseIterator,
+    MockResponseIterator,
+)
 from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
 from litellm.llms.base_llm.chat.transformation import BaseConfig
 from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
@@ -783,6 +786,8 @@ async def make_async_call_stream_helper(
             completion_stream = provider_config.get_model_response_iterator(
                 streaming_response=response.aiter_lines(), sync_stream=False
             )
+            if isinstance(completion_stream, BaseModelResponseIterator):
+                completion_stream.http_response = response
         # LOGGING
         logging_obj.post_call(
             input=messages,

litellm/proxy/common_request_processing.py

@@ -15,10 +15,12 @@
     Union,
 )
 
+import anyio
 import httpx
 import orjson
 from fastapi import HTTPException, Request, status
 from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.types import Receive, Scope, Send
 
 import litellm
 from litellm._logging import verbose_proxy_logger
@@ -150,6 +152,64 @@ def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict:
     return default_error
 
 
+async def _aclose_upstream_response(response: Any) -> None:
+    """Release the upstream HTTP connection when a stream ends for any
+    reason, including client disconnect. Mirrors the finally block of
+    async_data_generator in proxy_server.py."""
+    with anyio.CancelScope(shield=True):
+        if hasattr(response, "aclose"):
+            try:
+                await response.aclose()
+            except BaseException as e:
+                verbose_proxy_logger.debug(
+                    "error closing upstream response stream: %s", e
+                )
+
+
+class _UpstreamClosingStreamingResponse(StreamingResponse):
+    """StreamingResponse that always closes its body iterator and the wrapped
+    upstream generator.
+
+    When the client disconnects mid-stream, Starlette abandons the body
+    iterator without calling aclose(), leaving the upstream LLM connection
+    open until garbage collection; the backend (e.g. vLLM) keeps generating
+    into a dead pipe. The upstream generator is closed directly (not via the
+    body iterator) because aclose() on a never-started generator skips its
+    body, so a cascade through it would be a no-op if the client disconnects
+    before the first chunk is sent.
+    """
+
+    def __init__(
+        self,
+        content: AsyncGenerator[str, None],
+        *,
+        media_type: Optional[str] = None,
+        headers: Optional[dict] = None,
+        status_code: int = status.HTTP_200_OK,
+        upstream_generator: Optional[AsyncGenerator[str, None]] = None,
+    ) -> None:
+        super().__init__(
+            content, status_code=status_code, headers=headers, media_type=media_type
+        )
+        self._upstream_generator = upstream_generator
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        try:
+            await super().__call__(scope, receive, send)
+        finally:
+            with anyio.CancelScope(shield=True):
+                for target in (self.body_iterator, self._upstream_generator):
+                    aclose = getattr(target, "aclose", None)
+                    if aclose is None:
+                        continue
+                    try:
+                        await aclose()
+                    except BaseException as e:
+                        verbose_proxy_logger.debug(
+                            "error closing streaming generator: %s", e
+                        )
+
+
 async def create_response(
     generator: AsyncGenerator[str, None],
     media_type: str,
@@ -246,11 +306,12 @@ async def combined_generator() -> AsyncGenerator[str, None]:
             with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE):
                 yield chunk
 
-    return StreamingResponse(
+    return _UpstreamClosingStreamingResponse(
         combined_generator(),
         media_type=media_type,
         headers=headers,
         status_code=final_status_code,
+        upstream_generator=generator,
     )
 
 
@@ -1666,7 +1727,7 @@ def return_sse_chunk(chunk: Any) -> str:
             return chunk
 
     @staticmethod
-    async def async_streaming_data_generator(
+    async def async_streaming_data_generator(  # noqa: PLR0915
         response: Any,
         user_api_key_dict: UserAPIKeyAuth,
         request_data: dict,
@@ -1771,6 +1832,8 @@ async def async_streaming_data_generator(
                 code=getattr(e, "status_code", 500),
             )
             yield serialize_error(proxy_exception)
+        finally:
+            await _aclose_upstream_response(response)
 
     @staticmethod
     async def async_sse_data_generator(

tests/test_litellm/llms/base_llm/test_base_model_iterator.py

@@ -223,3 +223,38 @@ async def async_gen():
 
     assert len(chunks) == 1
     assert "response.created" in chunks[0]["text"]
+
+
+@pytest.mark.asyncio
+async def test_aclose_closes_attached_http_response():
+    """Regression for BerriAI/litellm#30244: CustomStreamWrapper.aclose() can
+    only release the upstream provider connection if the iterator exposes
+    aclose() and it reaches the underlying HTTP response. Without this, a
+    client disconnect leaves backends like vLLM generating into a dead pipe."""
+    from unittest.mock import AsyncMock, MagicMock
+
+    async def async_gen():
+        yield "data: {}"
+
+    iterator = BaseModelResponseIterator(
+        streaming_response=async_gen(), sync_stream=False
+    )
+    http_response = MagicMock()
+    http_response.aclose = AsyncMock()
+    iterator.http_response = http_response
+
+    await iterator.aclose()
+
+    http_response.aclose.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_aclose_is_noop_without_http_response():
+    async def async_gen():
+        yield "data: {}"
+
+    iterator = BaseModelResponseIterator(
+        streaming_response=async_gen(), sync_stream=False
+    )
+
+    await iterator.aclose()

tests/test_litellm/proxy/test_common_request_processing.py

@@ -1,3 +1,4 @@
+import asyncio
 import copy
 import datetime
 from typing import AsyncGenerator
@@ -19,6 +20,7 @@
     _is_azure_model_router_request,
     _override_openai_response_model,
     _parse_event_data_for_error,
+    _UpstreamClosingStreamingResponse,
     create_response,
 )
 from litellm.proxy.dd_span_tagger import DDSpanTagger
@@ -1853,3 +1855,182 @@ def test_depth_limit_prevents_infinite_loop(self):
         exc_a.__context__ = exc_b
         exc_b.__context__ = exc_a  # circular
         assert _has_attribute_error_in_chain(exc_a) is False
+
+
+class TestStreamCloseOnDisconnect:
+    """
+    Coverage for closing the upstream LLM stream when the client disconnects
+    mid-stream. Starlette abandons the response body iterator without calling
+    aclose(), so without these hooks the proxy->backend connection stays open
+    and the backend (e.g. vLLM) keeps generating into a dead pipe.
+    """
+
+    async def test_response_closes_body_iterator_when_task_cancelled(self):
+        """Cancellation landing in send() leaves the generator suspended at a
+        yield; only the response-level finally can close it."""
+        closed = asyncio.Event()
+
+        async def body():
+            try:
+                while True:
+                    yield "data: x\n\n"
+            finally:
+                closed.set()
+
+        response = _UpstreamClosingStreamingResponse(
+            body(), media_type="text/event-stream"
+        )
+
+        async def receive():
+            await asyncio.Event().wait()
+
+        async def send(message):
+            if message["type"] == "http.response.body":
+                await asyncio.Event().wait()
+
+        task = asyncio.create_task(response({"type": "http"}, receive, send))
+        await asyncio.sleep(0.05)
+        assert not closed.is_set()
+
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert closed.is_set()
+
+    async def test_response_closes_body_iterator_on_http_disconnect(self):
+        closed = asyncio.Event()
+        disconnected = asyncio.Event()
+        body_sends = 0
+
+        async def body():
+            try:
+                for i in range(1000):
+                    yield f"data: {i}\n\n"
+            finally:
+                closed.set()
+
+        response = _UpstreamClosingStreamingResponse(
+            body(), media_type="text/event-stream"
+        )
+
+        async def receive():
+            await disconnected.wait()
+            return {"type": "http.disconnect"}
+
+        async def send(message):
+            nonlocal body_sends
+            if message["type"] == "http.response.body":
+                body_sends += 1
+                if body_sends == 3:
+                    disconnected.set()
+                await asyncio.sleep(0.05)
+
+        await response({"type": "http"}, receive, send)
+
+        assert closed.is_set()
+        assert body_sends < 1000
+
+    async def test_upstream_closed_even_if_body_iterator_aclose_raises(self):
+        """A BaseException from body_iterator.aclose() (e.g. CancelledError)
+        must not prevent the upstream generator from being closed."""
+        upstream_closed = asyncio.Event()
+
+        class ExplodingIterator:
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                raise StopAsyncIteration
+
+            async def aclose(self):
+                raise asyncio.CancelledError()
+
+        async def upstream():
+            try:
+                yield "data: a\n\n"
+            finally:
+                upstream_closed.set()
+
+        upstream_gen = upstream()
+        await upstream_gen.__anext__()
+        response = _UpstreamClosingStreamingResponse(
+            ExplodingIterator(),
+            media_type="text/event-stream",
+            upstream_generator=upstream_gen,
+        )
+
+        async def receive():
+            await asyncio.Event().wait()
+
+        async def send(message):
+            pass
+
+        await response({"type": "http"}, receive, send)
+
+        assert upstream_closed.is_set()
+
+    async def test_create_response_closes_wrapped_generator_on_cancellation(self):
+        """End to end through create_response: the upstream-facing generator
+        must be closed even when the body iterator was never started (client
+        gone before the first chunk could be sent)."""
+        inner_closed = asyncio.Event()
+
+        async def wrapped():
+            try:
+                while True:
+                    yield "data: a\n\n"
+            finally:
+                inner_closed.set()
+
+        response = await create_response(
+            generator=wrapped(), media_type="text/event-stream", headers={}
+        )
+
+        async def receive():
+            await asyncio.Event().wait()
+
+        async def send(message):
+            await asyncio.Event().wait()
+
+        task = asyncio.create_task(response({"type": "http"}, receive, send))
+        await asyncio.sleep(0.05)
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert inner_closed.is_set()
+
+    async def test_async_streaming_data_generator_closes_upstream_on_early_close(
+        self,
+    ):
+        class FakeUpstream:
+            def __init__(self):
+                self.aclosed = False
+
+            def __aiter__(self):
+                return self
+
+            async def __anext__(self):
+                return {"type": "chunk"}
+
+            async def aclose(self):
+                self.aclosed = True
+
+        upstream = FakeUpstream()
+        gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
+            response=upstream,
+            user_api_key_dict=MagicMock(spec=UserAPIKeyAuth),
+            request_data={"model": "mock-model"},
+            proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()),
+            serialize_chunk=lambda c: "data: x\n\n",
+            serialize_error=lambda e: "data: error\n\n",
+        )
+
+        await gen.__anext__()
+        await gen.__anext__()
+        assert not upstream.aclosed
+
+        await gen.aclose()
+
+        assert upstream.aclosed