[Feat] Cancel upstream LLM request on client disconnect

litellm/constants.py

@@ -1287,6 +1287,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,
@@ -486,6 +487,39 @@ 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)
+        message = await request.receive()
+        if message.get("type") == "http.disconnect":
+            # 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
@@ -866,9 +900,11 @@ def _debug_log_request_payload(self) -> None:
                 "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Keys: %s",
                 len(_payload_str),
                 MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
-                list(self.data.keys())
-                if isinstance(self.data, dict)
-                else type(self.data).__name__,
+                (
+                    list(self.data.keys())
+                    if isinstance(self.data, dict)
+                    else type(self.data).__name__
+                ),
             )
         else:
             verbose_proxy_logger.debug(
@@ -1057,12 +1093,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
+            disconnect_task.cancel()
+        except asyncio.CancelledError:
+            disconnect_task.cancel()
+            if disconnect_event.is_set():
+                raise HTTPException(
+                    status_code=499,
+                    detail="Client disconnected the request",
+                )
+            raise
 
         response = responses[1]
 
@@ -1128,9 +1183,9 @@ async def base_process_llm_request(  # noqa: PLR0915
                 # aliasing/routing, but the OpenAI-compatible response `model` field should reflect
                 # what the client sent.
                 if requested_model_from_client:
-                    self.data[
-                        "_litellm_client_requested_model"
-                    ] = requested_model_from_client
+                    self.data["_litellm_client_requested_model"] = (
+                        requested_model_from_client
+                    )
 
                 # Streaming: attach a closure that fires after all guardrail
                 # end-of-stream blocks complete.  CSW.__anext__ stores the
@@ -1567,9 +1622,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,
@@ -1731,7 +1790,9 @@ async def async_streaming_data_generator(
         verbose_proxy_logger.debug("inside generator")
         try:
             str_so_far = ""
-            async for chunk in proxy_logging_obj.async_post_call_streaming_iterator_hook(
+            async for (
+                chunk
+            ) in proxy_logging_obj.async_post_call_streaming_iterator_hook(
                 user_api_key_dict=user_api_key_dict,
                 response=response,
                 request_data=request_data,
@@ -1959,9 +2020,9 @@ def _inject_cost_into_usage_dict(obj: dict, model_name: str) -> Optional[dict]:
 
             # Add cache-related fields to **params (handled by Usage.__init__)
             if cache_creation_input_tokens is not None:
-                usage_kwargs[
-                    "cache_creation_input_tokens"
-                ] = cache_creation_input_tokens
+                usage_kwargs["cache_creation_input_tokens"] = (
+                    cache_creation_input_tokens
+                )
             if cache_read_input_tokens is not None:
                 usage_kwargs["cache_read_input_tokens"] = cache_read_input_tokens
 

litellm/proxy/proxy_server.py

@@ -1608,34 +1608,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,56 @@
+"""
+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():
+    """Test that _check_request_disconnection cancels task and sets event when client disconnects."""
+    mock_request = AsyncMock()
+    mock_request.receive.side_effect = [
+        {"type": "http.request"},  # First call
+        {"type": "http.disconnect"},  # Second call - disconnect
+    ]
+
+    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()
+
+
+@pytest.mark.asyncio
+async def test_check_request_disconnection_no_disconnect():
+    """Test that _check_request_disconnection does not cancel task during normal operation."""
+    mock_request = AsyncMock()
+    mock_request.receive.return_value = {"type": "http.request"}
+
+    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)  # Let it run briefly
+    task.cancel()
+
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+
+    mock_llm_task.cancel.assert_not_called()
+    assert not disconnect_event.is_set()