From a4ce0f2ced0aca420f5d7ec4339d19010df2c060 Mon Sep 17 00:00:00 2001 From: vivekvar-dl Date: Sat, 8 Aug 2026 13:47:50 +0530 Subject: [PATCH] fix(proxy): deterministically close upstream streams for /v1/responses and /v1/messages The proxy's streaming cleanup closes the upstream response only when the stream object exposes aclose. LiteLLMCompletionStreamingIterator and BaseResponsesAPIStreamingIterator had none, and the /v1/messages passthrough generator never closed its httpx response, so releasing the provider connection depended on GC timing instead of an explicit close. The chat completions route got the explicit contract in #30245; this extends it to the two remaining streaming routes BaseResponsesAPIStreamingIterator gains a null-safe aclose (some subclasses bypass its constructor and carry no response), the bridge iterator delegates to its CustomStreamWrapper, and chunk_processor closes the upstream response in its finally after spend logging is scheduled Fixes #36123 --- .../streaming_handler.py | 4 + .../streaming_iterator.py | 3 + litellm/responses/streaming_iterator.py | 6 + .../test_streaming_handler.py | 114 ++++++++++++++++++ ...t_litellm_completion_streaming_iterator.py | 44 +++++++ .../responses/test_streaming_iterator.py | 47 ++++++++ 6 files changed, 218 insertions(+) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py create mode 100644 tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_streaming_iterator.py diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ff1c12d08d7..f77fcb98fbd 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -117,6 +117,10 @@ class PassThroughStreamingHandler: ) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) + try: + await response.aclose() + except Exception as e: + verbose_proxy_logger.debug("Error closing passthrough upstream response: %s", e) @staticmethod async def _route_streaming_logging_to_handler( diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index ddd05075763..1a289fc4444 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -106,6 +106,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._accumulated_provider_specific_fields: dict[str, Any] = {} self._custom_tool_names: set[str] = extract_custom_tool_names(self.responses_api_request.get("tools")) + async def aclose(self) -> None: + await self.litellm_custom_stream_wrapper.aclose() + def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing: Final = self._tool_output_index_by_call_id.get(call_id) if existing is not None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..fbce8444455 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -147,6 +147,12 @@ class BaseResponsesAPIStreamingIterator: This class contains shared logic for both synchronous and asynchronous iterators. """ + response: httpx.Response | None = None + + async def aclose(self) -> None: + if self.response is not None: + await self.response.aclose() + def __init__( self, response: httpx.Response, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py new file mode 100644 index 00000000000..7648491e81c --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -0,0 +1,114 @@ +"""Regression tests: chunk_processor must close the upstream httpx.Response +when the stream ends, including on client disconnect (generator aclose), so +the provider connection is released and backends like vLLM stop generating. +""" + +from collections.abc import AsyncIterator +from datetime import datetime, timezone +from typing import Final +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamHiddenParams, + AnthropicMessagesStreamingResponse, +) +from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + +class _TrackingStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + self.aclose_called = False + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + async def aclose(self) -> None: + self.aclose_called = True + + +def _streaming_response(chunks: list[bytes]) -> tuple[httpx.Response, _TrackingStream]: + stream: Final = _TrackingStream(chunks) + response: Final = httpx.Response( + 200, + stream=stream, + request=httpx.Request("POST", "http://upstream.test/v1/messages"), + ) + return response, stream + + +def _logging_obj_stub() -> MagicMock: + logging_obj: Final = MagicMock() + logging_obj.completion_start_time = None + logging_obj.model_call_details = {} + return logging_obj + + +def _chunk_processor(response: httpx.Response): + return PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-sonnet-4-5"}, + litellm_logging_obj=_logging_obj_stub(), + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(timezone.utc), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + + +def _drop_scheduled_coroutine(async_coroutine) -> None: + async_coroutine.close() + + +@pytest.mark.asyncio +async def test_chunk_processor_closes_response_on_client_disconnect(): + response, stream = _streaming_response([b"event: message_start\n\n", b"event: message_stop\n\n"]) + generator: Final = _chunk_processor(response) + + with patch("litellm.proxy.pass_through_endpoints.streaming_handler.GLOBAL_LOGGING_WORKER") as worker: + worker.ensure_initialized_and_enqueue.side_effect = _drop_scheduled_coroutine + first_chunk: Final = await generator.__anext__() + assert first_chunk == b"event: message_start\n\n" + await generator.aclose() + + assert stream.aclose_called + assert response.is_closed + + +@pytest.mark.asyncio +async def test_chunk_processor_closes_response_on_natural_completion(): + response, stream = _streaming_response([b"event: message_start\n\n", b"event: message_stop\n\n"]) + generator: Final = _chunk_processor(response) + + with patch("litellm.proxy.pass_through_endpoints.streaming_handler.GLOBAL_LOGGING_WORKER") as worker: + worker.ensure_initialized_and_enqueue.side_effect = _drop_scheduled_coroutine + collected: Final = [chunk async for chunk in generator] + + assert collected == [b"event: message_start\n\n", b"event: message_stop\n\n"] + assert stream.aclose_called + assert response.is_closed + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_response_aclose_closes_upstream_response(): + response, stream = _streaming_response([b"event: message_start\n\n"]) + wrapped: Final = AnthropicMessagesStreamingResponse( + completion_stream=_chunk_processor(response), + hidden_params=AnthropicMessagesStreamHiddenParams(additional_headers={}), + ) + + with patch("litellm.proxy.pass_through_endpoints.streaming_handler.GLOBAL_LOGGING_WORKER") as worker: + worker.ensure_initialized_and_enqueue.side_effect = _drop_scheduled_coroutine + first_chunk: Final = await wrapped.__anext__() + assert first_chunk == b"event: message_start\n\n" + await wrapped.aclose() + + assert stream.aclose_called + assert response.is_closed diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_streaming_iterator.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_streaming_iterator.py new file mode 100644 index 00000000000..ed8af6323a7 --- /dev/null +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_streaming_iterator.py @@ -0,0 +1,44 @@ +"""Regression test: the /v1/responses bridge iterator must expose aclose() +so the proxy's streaming cleanup can release the upstream provider +connection on client disconnect. +""" + +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, +) + + +class _TrackingCompletionStream: + def __init__(self): + self.aclose_called = False + + async def aclose(self) -> None: + self.aclose_called = True + + +@pytest.mark.asyncio +async def test_aclose_closes_wrapped_completion_stream(): + completion_stream: Final = _TrackingCompletionStream() + wrapper: Final = CustomStreamWrapper( + completion_stream=completion_stream, + model="gpt-4o-mini", + logging_obj=MagicMock(), + custom_llm_provider="openai", + ) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="gpt-4o-mini", + litellm_custom_stream_wrapper=wrapper, + request_input="hello", + responses_api_request={}, + ) + + await iterator.aclose() + + assert completion_stream.aclose_called + assert wrapper.completion_stream is None diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index 5b0f40fdf27..b840728c501 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -235,3 +235,50 @@ def test_sync_transport_error_before_completed_event_raises(): with pytest.raises(httpx.ReadError): for _ in iterator: pass + + +class _TrackingByteStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + self.aclose_called = False + + async def __aiter__(self): + for chunk in self._chunks: + yield chunk + + async def aclose(self) -> None: + self.aclose_called = True + + +@pytest.mark.asyncio +async def test_aclose_closes_underlying_httpx_response(): + stream = _TrackingByteStream([_sse_event({"type": "response.output_text.delta", "delta": "hi"})]) + response = httpx.Response( + 200, + stream=stream, + request=httpx.Request("POST", "http://upstream.test/v1/responses"), + ) + + iterator = ResponsesAPIStreamingIterator( + response=response, + model="gpt-4o-mini", + responses_api_provider_config=_mock_config(), + logging_obj=_logging_obj_stub(), + litellm_metadata={}, + custom_llm_provider="openai", + ) + + await iterator.__anext__() + await iterator.aclose() + + assert stream.aclose_called + assert response.is_closed + + +@pytest.mark.asyncio +async def test_aclose_is_a_noop_for_iterators_without_a_response(): + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + + iterator = BaseResponsesAPIStreamingIterator.__new__(BaseResponsesAPIStreamingIterator) + + await iterator.aclose()