diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 47ba75e79a4..d2bbfd38019 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1147,21 +1147,20 @@ async def _aclose_late_response(produced: Response) -> None: Closing an already-closed async generator is a no-op, so this is safe to call from both the relay's own teardown and the outer one. """ - if not isinstance(produced, StreamingResponse): - return - targets: Final = ( - (produced.body_iterator, produced.upstream_generator) - if isinstance(produced, _UpstreamClosingStreamingResponse) - else (produced.body_iterator,) - ) - for target in targets: - aclose = getattr(target, "aclose", None) - if aclose is None: - continue - try: - await aclose() - except BaseException as exc: # noqa: BLE001 # teardown must not mask why the stream ended - verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) + if isinstance(produced, StreamingResponse): + targets: Final = ( + (produced.body_iterator, produced.upstream_generator) + if isinstance(produced, _UpstreamClosingStreamingResponse) + else (produced.body_iterator,) + ) + for target in targets: + aclose = getattr(target, "aclose", None) + if aclose is None: + continue + try: + await aclose() + except BaseException as exc: # noqa: BLE001 # teardown must not mask why the stream ended + verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) if produced.background is not None: try: await produced.background() diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index a3366ed8990..f5f7f7e2a26 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -10245,3 +10245,20 @@ async def test_aclose_late_response_runs_background_task(): produced: Final = StreamingResponse(body(), background=BackgroundTask(mark)) await _aclose_late_response(produced) assert ran == [True] + + +@pytest.mark.asyncio +async def test_aclose_late_response_runs_background_task_for_non_streaming_response(): + from starlette.background import BackgroundTask + from starlette.responses import Response as StarletteResponse + + from litellm.proxy.common_request_processing import _aclose_late_response + + ran: list[bool] = [] + + async def mark() -> None: + ran.append(True) + + produced: Final = StarletteResponse(content=b"{}", background=BackgroundTask(mark)) + await _aclose_late_response(produced) + assert ran == [True]