ITADN

[Bug]: upstream LLM stream not closed on client disconnect for /v1/responses and /v1/messages

#36123Openxhejtman 创建于 14 天前
bug
X
xhejtmancommented
### Check for existing issues - [x] I have searched the existing issues and checked that my issue is not a duplicate. ### What happened? ### Follow-up to #30244 / #30245. That fix covers `/chat/completions` — `BaseModelResponseIterator` gained `self.http_response` plus an `aclose()`, attached at `llm_http_handler.py:836`. The other two streaming routes never reach it, so on v1.95.0 the backend still keeps generating after the client goes away. ### Why `_finalize_streaming_generator_cleanup` ends with `if hasattr(response, "aclose")`. That check fails on both remaining routes: **/v1/responses** — `response` is `LiteLLMCompletionStreamingIterator`. It holds the `CustomStreamWrapper` (which *does* have `aclose`) as `self.litellm_custom_stream_wrapper`, but defines no `aclose` itself, and neither does `BaseResponsesAPIStreamingIterator`. Nothing is closed. **/v1/messages** — `AnthropicMessagesStreamingResponse.aclose` exists and chains to `aclose_if_supported(self.completion_stream)`, but that stream is `PassThroughStreamingHandler.chunk_processor`, whose `finally` only schedules spend logging. `response.aclose()` is never called. ### Why closing the iterator isn't enough `httpx.Response.aiter_raw` calls `self.aclose()` only after natural exhaustion — there is no `finally`. httpcore does close the connection on `GeneratorExit`, but litellm retains both the `httpx.Response` and the `aiter_lines()` generator for the whole request, so neither is ever finalized. The socket stays open, the backend sees no disconnect, and vLLM/sglang keeps decoding into it. ### Steps to Reproduce Cancel request to `/v1/messages` or `/v1/responses` and see, if upstream stopped generation. ### Fix 57 lines across three files: `aclose()` on `BaseResponsesAPIStreamingIterator` closing `self.response` (also covers native `/v1/responses`), an override on `LiteLLMCompletionStreamingIterator` delegating to the wrapper, and `await response.aclose()` in `chunk_processor`'s existing `finally`. Both new closes run under `anyio.CancelScope(shield=True)` so they survive the cancellation that triggered the cleanup. ```diff diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 4dc1e0e..ed94ffa 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,6 +1,7 @@ from datetime import datetime from typing import List, Optional, Tuple +import anyio import httpx import litellm @@ -93,6 +94,19 @@ class PassThroughStreamingHandler: verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}") raise finally: + # Release the upstream connection before anything else, so a backend + # like vLLM/sglang sees the disconnect and aborts generation instead + # of decoding into a socket nobody reads. httpx auto-closes only on + # natural stream exhaustion -- not on GeneratorExit / cancellation -- + # so without this the connection lingers until GC. + if not response.is_closed: + with anyio.CancelScope(shield=True): + try: + await response.aclose() + except BaseException as e: # noqa: BLE001 + verbose_proxy_logger.debug( + "chunk_processor: error closing upstream response: %s", e + ) # GeneratorExit (raised on client disconnect) is not caught by # `except Exception`; the finally block ensures partial usage # still gets logged for spend tracking. See LIT-2642. diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index cf69654..07dad0b 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -106,6 +106,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._accumulated_provider_specific_fields: dict[str, Any] = {} self._custom_tool_names: set[str] = extract_custom_tool_names(self.responses_api_request.get("tools")) + async def aclose(self) -> None: + """Close the wrapped chat/completions stream so the upstream provider + connection is released when a bridged /v1/responses stream is abandoned. + + This class does not call ``BaseResponsesAPIStreamingIterator.__init__``, + so it has no ``self.response``; the httpx response lives inside the + ``CustomStreamWrapper``, whose own ``aclose`` already shields against + cancellation and swallows close errors. + """ + await self.litellm_custom_stream_wrapper.aclose() + def _get_or_assign_tool_output_index(self, call_id: str) -> int: existing = self._tool_output_index_by_call_id.get(call_id) if existing is not None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 357a7ec..803b0c1 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -9,6 +9,7 @@ from datetime import datetime from functools import lru_cache from typing import Any, Dict, List, Literal, Optional +import anyio import httpx from openai._streaming import SSEDecoder @@ -142,6 +143,37 @@ class BaseResponsesAPIStreamingIterator: self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + async def aclose(self) -> None: + """Close the upstream HTTP response so the provider connection is + released -- and a backend like vLLM/sglang aborts generation -- when the + stream is abandoned before its natural end (e.g. the client disconnected + mid-stream and the proxy is unwinding). + + httpx only closes the response when the byte stream is exhausted + naturally; on ``GeneratorExit`` / cancellation it does not, so the + connection would otherwise stay open and the backend would keep + decoding into a socket nobody reads. + + Mirrors ``BaseModelResponseIterator.aclose`` on the chat/completions + path. Subclasses that do not call this ``__init__`` (and therefore have + no ``self.response``) override this instead. + """ + response = getattr(self, "response", None) + if not isinstance(response, httpx.Response) or response.is_closed: + return + # Shield from the in-flight cancellation that triggered this cleanup so + # aclose() runs to completion and actually releases the connection. + with anyio.CancelScope(shield=True): + try: + await response.aclose() + except BaseException as e: # noqa: BLE001 + from litellm._logging import verbose_logger + + verbose_logger.debug( + "BaseResponsesAPIStreamingIterator.aclose: error closing upstream response: %s", + e, + ) + def _check_max_streaming_duration(self) -> None: """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: ``` ### What part of LiteLLM is this about? Proxy ### What LiteLLM version are you on ? v1.95.0 ### Twitter / LinkedIn details _No response_
0 条评论