fix(proxy): run late-relay background tasks for non-streaming responses too

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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]