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()