[WIP] Litellm client disconnect relay. Routing OSS PR into CI

litellm/constants.py

@@ -1309,6 +1309,9 @@
 DEFAULT_SOFT_BUDGET = float(
     os.getenv("DEFAULT_SOFT_BUDGET", 50.0)
 )  # by default all litellm proxy keys have a soft budget of 50.0
+DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS = int(
+    os.getenv("DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS", 600)
+)  # 10 minutes timeout for client disconnect checking in proxy
 # makes it clear this is a rate limit error for a litellm virtual key
 RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY = "LiteLLM Virtual Key user_api_key_hash"
 

litellm/proxy/common_request_processing.py

@@ -26,6 +26,7 @@
 from litellm._uuid import uuid
 from litellm.constants import (
     DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE,
+    DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS,
     DEFAULT_MAX_RECURSE_DEPTH,
     LITELLM_DETAILED_TIMING,
     MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
@@ -535,6 +536,38 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
     return False
 
 
+async def _check_request_disconnection(
+    request: Request,
+    llm_api_call_task,
+    disconnect_event: asyncio.Event,
+):
+    """
+    Asynchronously checks if the request is disconnected at regular intervals.
+    If the request is disconnected:
+    - cancels the litellm.router task (effective for non-streaming requests)
+    - sets disconnect_event so the caller can distinguish a client disconnect
+      from other sources of CancelledError (e.g. server shutdown)
+
+    Parameters:
+    - request: The request object to check for disconnection.
+    - llm_api_call_task: The asyncio gather future to cancel on disconnect.
+    - disconnect_event: Event set when the client disconnects.
+    Returns:
+    - None
+    """
+
+    # only run this function for configured timeout -> if these don't get cancelled -> we don't want the server to have many while loops
+    start_time = time.time()
+    while time.time() - start_time < DEFAULT_CLIENT_DISCONNECT_CHECK_TIMEOUT_SECONDS:
+        await asyncio.sleep(1)
+        if await request.is_disconnected():
+            # cancel the LLM API Call task if any passed - this is passed from individual providers
+            # Example OpenAI, Azure, VertexAI etc
+            llm_api_call_task.cancel()
+            disconnect_event.set()
+            return
+
+
 class ProxyBaseLLMRequestProcessing:
     def __init__(self, data: dict):
         self.data = data
@@ -1185,12 +1218,31 @@ async def base_process_llm_request(  # noqa: PLR0915
         )
         tasks.append(llm_call)
 
-        # wait for call to end
         llm_responses = asyncio.gather(
             *tasks
         )  # run the moderation check in parallel to the actual llm api call
 
-        responses = await llm_responses
+        # Execute the task to detect disconnection
+        disconnect_event = asyncio.Event()
+        disconnect_task = asyncio.create_task(
+            _check_request_disconnection(request, llm_responses, disconnect_event)
+        )
+
+        try:
+            # wait for call to end
+            # Note: for streaming requests llm_responses resolves quickly once the
+            # upstream connection is established; the ASGI transport layer handles
+            # cancellation of the upstream when the client disconnects mid-stream.
+            responses = await llm_responses
+        except asyncio.CancelledError:
+            if disconnect_event.is_set():
+                raise HTTPException(
+                    status_code=499,
+                    detail="Client disconnected the request",
+                )
+            raise
+        finally:
+            disconnect_task.cancel()
 
         response = responses[1]
 
@@ -1727,9 +1779,13 @@ async def _handle_llm_api_exception(
         version: Optional[str] = None,
     ):
         """Raises ProxyException (OpenAI API compatible) if an exception is raised"""
-        verbose_proxy_logger.exception(
-            f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}"
-        )
+        if isinstance(e, HTTPException) and e.status_code == 499:
+            verbose_proxy_logger.info("Client disconnected the request (499)")
+            raise e
+        else:
+            verbose_proxy_logger.exception(
+                f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}"
+            )
         # Allow callbacks to transform the error response
         transformed_exception = await proxy_logging_obj.post_call_failure_hook(
             user_api_key_dict=user_api_key_dict,

litellm/proxy/proxy_server.py

@@ -1773,34 +1773,6 @@ async def root_redirect():
 ### logger ###
 
 
-async def check_request_disconnection(request: Request, llm_api_call_task):
-    """
-    Asynchronously checks if the request is disconnected at regular intervals.
-    If the request is disconnected
-    - cancel the litellm.router task
-    - raises an HTTPException with status code 499 and detail "Client disconnected the request".
-
-    Parameters:
-    - request: Request: The request object to check for disconnection.
-    Returns:
-    - None
-    """
-
-    # only run this function for 10 mins -> if these don't get cancelled -> we don't want the server to have many while loops
-    start_time = time.time()
-    while time.time() - start_time < 600:
-        await asyncio.sleep(1)
-        if await request.is_disconnected():
-            # cancel the LLM API Call task if any passed - this is passed from individual providers
-            # Example OpenAI, Azure, VertexAI etc
-            llm_api_call_task.cancel()
-
-            raise HTTPException(
-                status_code=499,
-                detail="Client disconnected the request",
-            )
-
-
 def _resolve_typed_dict_type(typ):
     """Resolve the actual TypedDict class from a potentially wrapped type."""
     from typing_extensions import _TypedDictMeta  # type: ignore

tests/test_litellm/proxy/test_client_disconnection.py

@@ -0,0 +1,55 @@
+"""
+Test client disconnection detection functionality.
+"""
+
+import asyncio
+
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from litellm.proxy.common_request_processing import _check_request_disconnection
+
+
+@pytest.mark.asyncio
+async def test_check_request_disconnection_with_disconnect():
+    """Disconnect path: polling sees disconnected only after is_disconnected becomes True."""
+    mock_request = MagicMock(spec=["is_disconnected"])
+    mock_request.is_disconnected = AsyncMock(side_effect=[False, True])
+
+    mock_llm_task = MagicMock()  # sync mock so .cancel() doesn't return a coroutine
+    disconnect_event = asyncio.Event()
+
+    with patch(
+        "litellm.proxy.common_request_processing.asyncio.sleep", new_callable=AsyncMock
+    ):
+        await _check_request_disconnection(
+            mock_request, mock_llm_task, disconnect_event
+        )
+
+    mock_llm_task.cancel.assert_called_once()
+    assert disconnect_event.is_set()
+    assert mock_request.is_disconnected.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_check_request_disconnection_no_disconnect():
+    """Cancel watcher mid-flight: LLM task must not be cancelled like a disconnect."""
+    mock_request = MagicMock(spec=["is_disconnected"])
+    mock_request.is_disconnected = AsyncMock(return_value=False)
+
+    mock_llm_task = MagicMock()  # sync mock so .cancel() doesn't return a coroutine
+    disconnect_event = asyncio.Event()
+
+    task = asyncio.create_task(
+        _check_request_disconnection(mock_request, mock_llm_task, disconnect_event)
+    )
+    await asyncio.sleep(0.1)
+    task.cancel()
+
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+
+    mock_llm_task.cancel.assert_not_called()
+    assert not disconnect_event.is_set()