diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 2ecc95b7b32..9969073e2ae 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,9 +5,11 @@ import traceback from datetime import datetime from typing import Any, Dict, List, Optional +import anyio import httpx import litellm +from litellm._logging import verbose_logger from litellm.constants import ( LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING, @@ -89,6 +91,37 @@ class BaseResponsesAPIStreamingIterator: self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + async def aclose(self) -> None: + """ + Release the underlying httpx.Response back to the connection pool. + + Safe to call multiple times; only the first call performs cleanup. + Shielded from anyio cancellation so cleanup awaits complete even when + the surrounding task is being cancelled (e.g. client disconnect). + + Mirrors CustomStreamWrapper.aclose (see PR #21213) for the Responses + API path. Fixes #26250 — without this, client disconnects leak the + upstream connection until the pool is exhausted. + """ + response = self.response + if response is None: + return + self.response = None # type: ignore[assignment] + self.finished = True + with anyio.CancelScope(shield=True): + try: + if hasattr(response, "aclose"): + await response.aclose() + elif hasattr(response, "close"): + result = response.close() + if result is not None: + await result + except BaseException as e: + verbose_logger.debug( + "BaseResponsesAPIStreamingIterator.aclose: error closing response: %s", + e, + ) + def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: @@ -771,7 +804,6 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): # WebSocket mode streaming (bidirectional forwarding) # --------------------------------------------------------------------------- -from litellm._logging import verbose_logger from litellm.litellm_core_utils.thread_pool_executor import executor as _ws_executor RESPONSES_WS_LOGGED_EVENT_TYPES = [ diff --git a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py index e9181d810e1..b387d751fc4 100644 --- a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py +++ b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py @@ -17,7 +17,7 @@ import os import sys from datetime import datetime from typing import Any, Dict, Optional -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, Mock, patch import pytest @@ -583,3 +583,82 @@ class TestBaseResponsesAPIStreamingIterator: mock_logging_obj.async_failure_handler.assert_not_called() mock_logging_obj.failure_handler.assert_not_called() + +class TestBaseResponsesAPIStreamingIteratorAclose: + """ + Tests for BaseResponsesAPIStreamingIterator.aclose() — ensures the + underlying httpx.Response is released back to the connection pool when a + stream is abandoned mid-iteration. See issue #26250 and PR #21213 for the + equivalent fix on the chat-completions path. + """ + + def _build_iterator(self, mock_response): + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_config = Mock(spec=BaseResponsesAPIConfig) + return BaseResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-4", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + ) + + @pytest.mark.asyncio + async def test_aclose_calls_response_aclose(self): + """aclose() awaits response.aclose(), nulls self.response, and sets finished.""" + mock_response = Mock() + mock_response.headers = {} + mock_response.aclose = AsyncMock() + + iterator = self._build_iterator(mock_response) + + await iterator.aclose() + + mock_response.aclose.assert_awaited_once() + assert iterator.response is None + assert iterator.finished is True + + @pytest.mark.asyncio + async def test_aclose_is_idempotent(self): + """Calling aclose() twice only closes the response once.""" + mock_response = Mock() + mock_response.headers = {} + mock_response.aclose = AsyncMock() + + iterator = self._build_iterator(mock_response) + + await iterator.aclose() + await iterator.aclose() + + mock_response.aclose.assert_awaited_once() + + @pytest.mark.asyncio + async def test_aclose_swallows_close_errors(self): + """Errors from the underlying close are logged, not raised.""" + mock_response = Mock() + mock_response.headers = {} + mock_response.aclose = AsyncMock(side_effect=RuntimeError("boom")) + + iterator = self._build_iterator(mock_response) + + # Must not raise — cleanup errors are swallowed so they don't mask + # the user's original exception path. + await iterator.aclose() + + mock_response.aclose.assert_awaited_once() + assert iterator.response is None + + @pytest.mark.asyncio + async def test_aclose_falls_back_to_close_when_no_aclose(self): + """If the response exposes only close (not aclose), aclose() invokes close.""" + mock_response = Mock(spec=["headers", "close"]) + mock_response.headers = {} + mock_response.close = Mock(return_value=None) + + iterator = self._build_iterator(mock_response) + + await iterator.aclose() + + mock_response.close.assert_called_once() + assert iterator.response is None +