diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 0eaa700ffd..0aeb83a64b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -7,6 +7,7 @@ import time import traceback from typing import Any, Callable, Dict, List, Optional, Union, cast +import anyio import httpx from pydantic import BaseModel @@ -160,8 +161,6 @@ class CustomStreamWrapper: # Shield from anyio cancellation so cleanup awaits can complete. # Without this, CancelledError is thrown into every await during # task group cancellation, preventing HTTP connection release. - import anyio - with anyio.CancelScope(shield=True): try: if hasattr(self.completion_stream, "aclose"): diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d0969de7ca..a02bc7f9e5 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,6 @@ import json import logging import traceback from datetime import datetime -from functools import partial from typing import ( TYPE_CHECKING, Any, @@ -15,13 +14,10 @@ from typing import ( Union, ) -import anyio import httpx import orjson from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse -from starlette._utils import collapse_excgroups -from starlette.types import Receive, Scope, Send import litellm from litellm._logging import verbose_proxy_logger @@ -142,39 +138,6 @@ def _extract_error_from_sse_chunk(event_line: Union[str, bytes]) -> dict: return default_error -async def _disconnect_aware_call( - self: StreamingResponse, scope: Scope, receive: Receive, send: Send -) -> None: - """Patched StreamingResponse.__call__ that detects client disconnects. - - Starlette >= 0.45.3 relies on ASGI spec 2.4 where send() should raise OSError - on client disconnect. Uvicorn does not implement this — send() silently returns. - This means async generators keep running forever after clients disconnect, - leaking upstream HTTP connections. - - This restores the pre-0.45.3 behavior: a concurrent task listens for - http.disconnect and cancels the streaming task group when detected. - - See: https://github.com/encode/starlette/pull/2732 - See: https://github.com/encode/uvicorn/pull/2276 - """ - with collapse_excgroups(): - async with anyio.create_task_group() as task_group: - - async def wrap(func: Callable[[], Any]) -> None: - await func() - task_group.cancel_scope.cancel() - - task_group.start_soon(wrap, partial(self.stream_response, send)) - await wrap(partial(self.listen_for_disconnect, receive)) - - if self.background is not None: - await self.background() - - -StreamingResponse.__call__ = _disconnect_aware_call # type: ignore[assignment] - - async def create_response( generator: AsyncGenerator[str, None], media_type: str, diff --git a/litellm/router.py b/litellm/router.py index f767ff59a5..0c55d93066 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1497,6 +1497,7 @@ class Router: return await self._async_generator.__anext__() async def stream_with_fallbacks(): + fallback_response = None # Track for cleanup in finally try: async for item in model_response: yield item @@ -1596,13 +1597,17 @@ class Router: ) raise fallback_error finally: - # Close the underlying stream to release the HTTP connection + # Close the underlying streams to release HTTP connections # back to the connection pool when the generator is closed # (e.g. on client disconnect). - # Shield from anyio cancellation so the await can complete. - if hasattr(model_response, "aclose"): - with anyio.CancelScope(shield=True): + # Shield from anyio cancellation so the awaits can complete. + with anyio.CancelScope(shield=True): + if hasattr(model_response, "aclose"): await model_response.aclose() + if fallback_response is not None and hasattr( + fallback_response, "aclose" + ): + await fallback_response.aclose() return FallbackStreamWrapper(stream_with_fallbacks()) diff --git a/pyproject.toml b/pyproject.toml index be15013267..a42bdafa1a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ pydantic = "^2.5.0" jsonschema = ">=4.23.0,<5.0.0" numpydoc = {version = "*", optional = true} # used in utils.py -uvicorn = {version = "^0.31.1", optional = true} +uvicorn = {version = ">=0.32.1", optional = true} uvloop = {version = "^0.21.0", optional = true, markers="sys_platform != 'win32'"} gunicorn = {version = "^23.0.0", optional = true} fastapi = {version = ">=0.120.1", optional = true}