From e3154e974284085117b62d18b3fc140f151b8f20 Mon Sep 17 00:00:00 2001 From: Vedant Agarwal Date: Tue, 30 Jun 2026 00:05:21 +0800 Subject: [PATCH] fix(langgraph): release httpx connection when SSE stream ends or is abandoned LangGraphSSEStreamIterator held the streaming httpx response but defined neither close() nor aclose(), so CustomStreamWrapper.aclose() (which dispatches to aclose()/close()) silently no-opped and the connection was never returned to the pool on client disconnect or early stream abandonment. __next__/__anext__ also swallowed mid-stream errors without releasing the response. Add close()/aclose() that release self.response, and close on the iterator's termination/error paths. Same connection-leak class as #30271. Fixes #31619 --- litellm/llms/langgraph/chat/sse_iterator.py | 24 ++++ .../test_streaming_connection_cleanup.py | 126 +++++++++++++++++- 2 files changed, 144 insertions(+), 6 deletions(-) diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 2eb17b4d4b4..6b9228e7ec3 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -177,6 +177,26 @@ 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).""" + if self.response is not None: + response = self.response + self.response = None + 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.""" + if self.response is not None: + response = self.response + self.response = None + await response.aclose() + def __next__(self) -> ModelResponseStream: """Sync iteration - parse SSE events and yield ModelResponse chunks.""" try: @@ -196,6 +216,7 @@ class LangGraphSSEStreamIterator: raise StopIteration except StopIteration: + self.close() raise except httpx.StreamConsumed: raise StopIteration @@ -203,6 +224,7 @@ class LangGraphSSEStreamIterator: raise StopIteration except Exception as e: verbose_logger.error(f"Error in LangGraph SSE stream: {str(e)}") + self.close() raise StopIteration async def __anext__(self) -> ModelResponseStream: @@ -224,6 +246,7 @@ class LangGraphSSEStreamIterator: raise StopAsyncIteration except StopAsyncIteration: + await self.aclose() raise except httpx.StreamConsumed: raise StopAsyncIteration @@ -231,4 +254,5 @@ class LangGraphSSEStreamIterator: raise StopAsyncIteration except Exception as e: verbose_logger.error(f"Error in LangGraph SSE stream: {str(e)}") + await self.aclose() raise StopAsyncIteration diff --git a/tests/test_litellm/test_streaming_connection_cleanup.py b/tests/test_litellm/test_streaming_connection_cleanup.py index 5a81a3ffb17..6b62feb89e8 100644 --- a/tests/test_litellm/test_streaming_connection_cleanup.py +++ b/tests/test_litellm/test_streaming_connection_cleanup.py @@ -58,9 +58,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) @@ -142,6 +140,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.""" @@ -230,9 +346,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