From 08c3bf61921d3d1543bddba1b7b7dce94bded799 Mon Sep 17 00:00:00 2001
From: yucheng <yucheng@berri.ai>
Date: Sat, 26 Sep 2026 02:36:01 +0000
Subject: [PATCH] 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    | 29 +++++++++----------
 .../proxy/test_common_request_processing.py   | 17 +++++++++++
 2 files changed, 31 insertions(+), 15 deletions(-)

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]