Merge b0c081abc2 into f2118fc767

litellm/llms/langgraph/chat/sse_iterator.py

@@ -172,6 +172,22 @@ class LangGraphSSEStreamIterator:
 
         return chunk
 
+    def close(self) -> None:
+        """Release the underlying streaming response so the connection is
+        returned to the pool when the stream ends or is abandoned (sync).
+        httpx makes repeated calls a no-op."""
+        self.response.close()
+
+    async def aclose(self) -> None:
+        """Release the underlying streaming response so the connection is
+        returned to the pool when the stream ends or is abandoned (async).
+
+        ``CustomStreamWrapper.aclose()`` dispatches to this on client
+        disconnect / early stream abandonment; without it the httpx
+        connection is never released and the pool leaks. httpx makes
+        repeated calls a no-op."""
+        await self.response.aclose()
+
     def __next__(self) -> ModelResponseStream:
         """Sync iteration - parse SSE events and yield ModelResponse chunks."""
         try:
@@ -191,6 +207,7 @@ class LangGraphSSEStreamIterator:
             raise StopIteration
 
         except StopIteration:
+            self.close()
             raise
         except httpx.StreamConsumed:
             raise StopIteration
@@ -198,6 +215,7 @@ class LangGraphSSEStreamIterator:
             raise StopIteration
         except Exception as e:
             verbose_logger.error("Error in LangGraph SSE stream: %s", e)
+            self.close()
             raise StopIteration
 
     async def __anext__(self) -> ModelResponseStream:
@@ -219,6 +237,7 @@ class LangGraphSSEStreamIterator:
             raise StopAsyncIteration
 
         except StopAsyncIteration:
+            await self.aclose()
             raise
         except httpx.StreamConsumed:
             raise StopAsyncIteration
@@ -226,4 +245,5 @@ class LangGraphSSEStreamIterator:
             raise StopAsyncIteration
         except Exception as e:
             verbose_logger.error("Error in LangGraph SSE stream: %s", e)
+            await self.aclose()
             raise StopAsyncIteration

tests/unit/test_streaming_connection_cleanup.py

@@ -55,9 +55,7 @@ async def test_aiohttp_transport_response_uses_stream_not_content():
             return Resp()
 
     transport = LiteLLMAiohttpTransport(client=lambda: FakeSession())  # type: ignore
-    response = await transport.handle_async_request(
-        httpx.Request("GET", "http://example.com")
-    )
+    response = await transport.handle_async_request(httpx.Request("GET", "http://example.com"))
 
     assert isinstance(response.stream, AiohttpResponseStream)
 
@@ -139,6 +137,124 @@ async def test_aclose_prefers_aclose_over_close():
     assert not close_called
 
 
+@pytest.mark.asyncio
+async def test_langgraph_sse_iterator_aclose_releases_response():
+    """LangGraphSSEStreamIterator must expose aclose() so CustomStreamWrapper.aclose()
+    releases the underlying httpx response; otherwise the connection leaks on a client
+    disconnect / abandoned stream."""
+    from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
+
+    response_released = False
+
+    class FakeResponse:
+        async def aclose(self):
+            nonlocal response_released
+            response_released = True
+
+    iterator = LangGraphSSEStreamIterator(response=FakeResponse(), model="x")  # type: ignore
+    wrapper = CustomStreamWrapper(
+        completion_stream=iterator,
+        model=None,
+        logging_obj=MagicMock(),
+        custom_llm_provider=None,
+    )
+
+    await wrapper.aclose()
+    assert response_released, (
+        "LangGraph streaming response was not released by CustomStreamWrapper.aclose() — the httpx connection leaks"
+    )
+
+
+_LANGGRAPH_AI_LINE = 'data: ["messages", [[{"type": "ai", "content": "hi"}]]]'
+
+
+def test_langgraph_iterator_close_on_sync_completion():
+    """Sync iteration to completion releases the underlying response."""
+    from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
+
+    closed = False
+
+    class FakeResponse:
+        def iter_lines(self):
+            return iter([_LANGGRAPH_AI_LINE])
+
+        def close(self):
+            nonlocal closed
+            closed = True
+
+    list(LangGraphSSEStreamIterator(response=FakeResponse(), model="x"))  # type: ignore
+    assert closed
+
+
+def test_langgraph_iterator_close_on_sync_error():
+    """A mid-stream error releases the underlying response (sync)."""
+    from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
+
+    closed = False
+
+    class FakeResponse:
+        def iter_lines(self):
+            def gen():
+                raise RuntimeError("boom")
+                yield  # pragma: no cover
+
+            return gen()
+
+        def close(self):
+            nonlocal closed
+            closed = True
+
+    list(LangGraphSSEStreamIterator(response=FakeResponse(), model="x"))  # type: ignore
+    assert closed
+
+
+@pytest.mark.asyncio
+async def test_langgraph_iterator_aclose_on_async_completion():
+    """Async iteration to completion releases the underlying response."""
+    from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
+
+    closed = False
+
+    class FakeResponse:
+        def aiter_lines(self):
+            async def gen():
+                yield _LANGGRAPH_AI_LINE
+
+            return gen()
+
+        async def aclose(self):
+            nonlocal closed
+            closed = True
+
+    async for _ in LangGraphSSEStreamIterator(response=FakeResponse(), model="x"):  # type: ignore
+        pass
+    assert closed
+
+
+@pytest.mark.asyncio
+async def test_langgraph_iterator_aclose_on_async_error():
+    """A mid-stream error releases the underlying response (async)."""
+    from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator
+
+    closed = False
+
+    class FakeResponse:
+        def aiter_lines(self):
+            async def gen():
+                raise RuntimeError("boom")
+                yield  # pragma: no cover
+
+            return gen()
+
+        async def aclose(self):
+            nonlocal closed
+            closed = True
+
+    async for _ in LangGraphSSEStreamIterator(response=FakeResponse(), model="x"):  # type: ignore
+        pass
+    assert closed
+
+
 @pytest.mark.asyncio
 async def test_aclose_completes_under_cancellation():
     """aclose() must shield cleanup from CancelledError so streams actually close."""
@@ -227,9 +343,7 @@ async def test_stream_with_fallbacks_closes_stream_on_generator_close():
         break
     await result.aclose()
 
-    assert (
-        stream_closed
-    ), "model_response stream was not closed by stream_with_fallbacks finally block"
+    assert stream_closed, "model_response stream was not closed by stream_with_fallbacks finally block"
 
 
 @pytest.mark.asyncio