From 5b31216caef4418caf739dcc103dcc41314cd3ff Mon Sep 17 00:00:00 2001
From: yryzhan <yryzhan@users.noreply.github.com>
Date: Wed, 20 May 2026 21:16:38 +0200
Subject: [PATCH] fix(streaming): cleanup _queue_wrapper in aclose to prevent
 thread leak
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

On client disconnect, aclose() is called but previously did not stop
the producer thread — it would block indefinitely on a full queue.
Now aclose() calls wrapper.close() first, setting the stop event
so the producer exits promptly.
---
 .../litellm_core_utils/streaming_handler.py   |  3 ++
 .../test_streaming_handler.py                 | 39 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index 721287a6f41..0844a5ef9c1 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -268,6 +268,9 @@ class CustomStreamWrapper:
         return self
 
     async def aclose(self):
+        if self._queue_wrapper is not None:
+            self._queue_wrapper.close()
+            self._queue_wrapper = None
         if self.completion_stream is not None:
             stream_to_close = self.completion_stream
             self.completion_stream = None
diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
index 58b0eb11eaa..7b81171905d 100644
--- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
+++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py
@@ -2353,6 +2353,45 @@ async def test_anext_str_completion_stream():
         pass
 
 
+@pytest.mark.asyncio
+async def test_queue_wrapper_cleanup_on_aclose():
+    """aclose() stops the producer thread to prevent leaks on client disconnect."""
+    loop = asyncio.get_running_loop()
+
+    def infinite_iter():
+        i = 0
+        while True:
+            i += 1
+            yield i
+
+    wrapper = _SyncIteratorToQueue(infinite_iter(), loop)
+    await wrapper.get()
+
+    from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
+
+    logging_obj = MagicMock()
+    logging_obj.model_call_details = {"litellm_params": {}}
+    logging_obj.stream_options = None
+    logging_obj.messages = []
+
+    stream = CustomStreamWrapper(
+        completion_stream=infinite_iter(),
+        model="test-model",
+        logging_obj=logging_obj,
+        custom_llm_provider="openai",
+    )
+
+    # Force _queue_wrapper to exist
+    stream._queue_wrapper = wrapper
+    assert wrapper._thread.is_alive()
+
+    await stream.aclose()
+
+    assert stream._queue_wrapper is None
+    await asyncio.sleep(0.6)
+    assert not wrapper._thread.is_alive()
+
+
 @pytest.mark.asyncio
 async def test_queue_wrapper_cleanup_on_httpx_timeout():
     """_queue_wrapper is cleaned up when httpx.TimeoutException is raised during streaming (lines 2306-2307)."""