diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index cf1fd6f786e..eb05c68f23c 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c9abaad908..82df782f97e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a61bce833ad..4ffb634b8ea 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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( diff --git a/tests/test_litellm/llms/base_llm/test_base_model_iterator.py b/tests/test_litellm/llms/base_llm/test_base_model_iterator.py index 7a001b26002..fb1ac669599 100644 --- a/tests/test_litellm/llms/base_llm/test_base_model_iterator.py +++ b/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() diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 781b559651f..caf2b2dc744 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/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